Fix flaky ICP test_heterogeneous_inputs

Summary:
`TestICP.test_heterogeneous_inputs` has been failing intermittently for a long
time. The seed was already bumped from 4 to 14 in D80625966 for the same
reason, which relocated the failure rather than removing it.

Root cause: the test aligned two independent random point clouds, and with
`estimate_scale=True` that problem is ill-posed. ICP almost always collapses
`X` onto a single point of `Y`, driving `s` to ~1e-16 and the rmse to zero.
That solution fits perfectly, but the rotation of a cloud that has shrunk to a
point is completely unconstrained, so the batched run and the per-cloud runs
each returned an arbitrary, and different, `R`. Over 600 seeds of the old data
~45% of batch elements collapsed, and every element that exceeded the
tolerance was a collapsed one, with `R` the only quantity that disagreed
(`T`, `s` and `Xt` always matched). This is the same non-uniqueness that
`corresponding_points_alignment` warns about with "Excessively low rank of
cross-correlation".

Which seeds tripped over it came down to float32 rounding: the batched path
sums over the zero-weighted padding and the per-cloud path does not, so the
two differ by ~1e-7, and that difference decides which arbitrary rotation
comes out. It therefore moves with GPU model, BLAS version and TF32, which is
why picking a seed was never a fix - forcing TF32 on makes seed 14 fail
immediately.

Fix: build `Y` as a rigidly moved copy of `X` plus a few extra points. The
clouds still have different sizes within a batch and between `X` and `Y`, so
the padding and masking path is exercised exactly as before, but the alignment
now has a well-determined optimum that ICP cannot collapse.

Also loosen `atol` from 1e-5 to 1e-4. That is needed independently of the
collapse: the two runs sum a different number of terms and so round
differently, and the legitimate deviation on `Xt` reaches 1.17e-5, above the
old tolerance. That was a second latent failure waiting to happen.

The test no longer emits the "Excessively low rank of cross-correlation"
warning, which is the collapse disappearing.

___

Differential Revision: D117539518

fbshipit-source-id: 73e3d70a8f7547de7179b71e6a705fc37f4920c7

tests/test_points_alignment.py

@@ -159,6 +159,38 @@ class TestICP(TestCaseMixin, unittest.TestCase):
                     self.assertClose(s_init, s, atol=atol)
                     self.assertClose(Xt_init, Xt, atol=atol)
 
+    @staticmethod
+    def _init_heterogeneous_icp_problem(batch_size, max_n_points, device):
+        """
+        Generate a batch of randomly-sized point clouds `X` together with
+        `Y`, a rigidly moved copy of each cloud with a few extra points
+        added, so that the clouds in `Y` have different sizes from those
+        in `X`.
+
+        `Y` deliberately covers the whole of `X`. Two independent random
+        clouds instead make the problem ill-posed: with `estimate_scale`,
+        ICP then usually collapses `X` onto a single point of `Y`, driving
+        the scale to ~1e-16 and the rmse to zero. That solution fits
+        perfectly but leaves the rotation completely unconstrained, so two
+        runs of the same algorithm may return wildly different rotations
+        and can't be compared with each other.
+        """
+        n_points = torch.randint(
+            low=6, high=max_n_points, size=(batch_size,), device=device
+        )
+        X_list = [torch.randn(int(n), 3, device=device) for n in n_points]
+        R = rotation_conversions.axis_angle_to_matrix(
+            0.15 * torch.randn(batch_size, 3, device=device)
+        )
+        T = 0.1 * torch.randn(batch_size, 3, device=device)
+        Y_list = []
+        for i, X in enumerate(X_list):
+            n_extra = int(torch.randint(low=1, high=5, size=(), device=device))
+            covered = torch.cat([X, torch.randn(n_extra, 3, device=device)], dim=0)
+            Y_list.append(covered @ R[i] + T[i] + 0.01 * torch.randn_like(covered))
+
+        return Pointclouds(X_list), Pointclouds(Y_list)
+
     def test_heterogeneous_inputs(self, batch_size=7):
         """
         Tests whether we get the same result when running ICP on
@@ -171,17 +203,9 @@ class TestICP(TestCaseMixin, unittest.TestCase):
         for estimate_scale in (True, False):
             for max_n_points in (10, 30, 100):
                 # initialize ground truth point clouds
-                X_pcl, Y_pcl = [
+                X_pcl, Y_pcl = self._init_heterogeneous_icp_problem(
-                    TestCorrespondingPointsAlignment.init_point_cloud(
+                    batch_size, max_n_points, device
-                        batch_size=batch_size,
+                )
-                        n_points=max_n_points,
-                        dim=3,
-                        device=device,
-                        use_pointclouds=True,
-                        random_pcl_size=True,
-                    )
-                    for _ in range(2)
-                ]
 
                 # get the padded versions and their num of points
                 X_padded = X_pcl.points_padded()
@@ -227,8 +251,11 @@ class TestICP(TestCaseMixin, unittest.TestCase):
                     torch.cat([x.RTs[i] for x in icp_results], dim=0) for i in range(3)
                 ]
 
-                # check that both sets of transforms are the same
+                # check that both sets of transforms are the same.
-                atol = 1e-5
+                # The two runs sum over a different number of points - the
+                # batched one includes the (zero-weighted) padding - so they
+                # round differently, by up to ~1e-5 on the transformed clouds.
+                atol = 1e-4
                 self.assertClose(R_pcl, R, atol=atol)
                 self.assertClose(T_pcl, T, atol=atol)
                 self.assertClose(s_pcl, s, atol=atol)