Added an option to flatten_dims in RobustCostFunction.

tests/core/test_robust_cost.py

@@ -136,3 +136,90 @@ def test_mask_jacobians(loss_cls):
     torch.testing.assert_close(err, err_expected)
     for j1, j2 in zip(jac, jac_expected):
         torch.testing.assert_close(j1, j2)
+
+
+def _data_model(a, b, x):
+    return a * x.square() + b
+
+
+def _generate_data(num_points=100, a=1, b=0.5, noise_factor=0.01):
+    data_x = torch.rand((1, num_points))
+    noise = torch.randn((1, num_points)) * noise_factor
+    return data_x, _data_model(a, b, data_x) + noise
+
+
+@pytest.mark.parametrize("batch_size", [1, 4])
+def test_flatten_dims(batch_size):
+    # This creates two objectives for a regression problem
+    # and compares their linearization
+    #   - Obj1: N 1d robust costs functions each evaluating one residual term
+    #   - Obj2: 1 Nd robust cost function that evaluates all residuals at once
+    # Data for regression problem y ~ Normal(Ax^2 + B, sigma)
+    n = 10
+    data_x, data_y = _generate_data(num_points=n)
+    data_y[:, :5] = 1000  # include some extreme outliers for robust cost
+
+    # optimization variables are of type Vector with 2 degrees of freedom (dof)
+    # one for A and one for B
+    ab = th.Vector(2, name="ab")
+
+    def residual_fn(optim_vars, aux_vars):
+        ab = optim_vars[0]
+        x, y = aux_vars
+        return y.tensor - _data_model(ab.tensor[:, :1], ab.tensor[:, 1:], x.tensor)
+
+    w = th.ScaleCostWeight(0.5)
+    log_loss_radius = th.as_variable(0.5)
+
+    # First create an objective with individual cost functions per error term
+    # Need individual aux variables to represent data of each residual terms
+    xs = [th.Vector(1, name=f"x{i}") for i in range(n)]
+    ys = [th.Vector(1, name=f"y{i}") for i in range(n)]
+    obj_unrolled = th.Objective()
+    for i in range(n):
+        obj_unrolled.add(
+            th.RobustCostFunction(
+                th.AutoDiffCostFunction(
+                    (ab,), residual_fn, 1, aux_vars=(xs[i], ys[i]), cost_weight=w
+                ),
+                th.HuberLoss,
+                log_loss_radius,
+                name=f"rcf{i}",
+            )
+        )
+    lin_unrolled = th.DenseLinearization(obj_unrolled)
+    th_inputs = {f"x{i}": data_x[:, i].view(1, 1) for i in range(data_x.shape[1])}
+    th_inputs.update({f"y{i}": data_y[:, i].view(1, 1) for i in range(data_y.shape[1])})
+    th_inputs.update({"ab": torch.rand((batch_size, 2))})
+    obj_unrolled.update(th_inputs)
+    lin_unrolled.linearize()
+
+    # Now one with a single vectorized cost function, and flatten_dims=True
+    # Residual terms call all be represented with "batched" data variables
+    xb = th.Vector(n, name="xb")
+    yb = th.Vector(n, name="yb")
+    obj_flattened = th.Objective()
+    obj_flattened.add(
+        th.RobustCostFunction(
+            th.AutoDiffCostFunction(
+                (ab,), residual_fn, n, aux_vars=(xb, yb), cost_weight=w
+            ),
+            th.HuberLoss,
+            log_loss_radius,
+            name="rcf",
+            flatten_dims=True,
+        )
+    )
+    lin_flattened = th.DenseLinearization(obj_flattened)
+    th_inputs = {
+        "xb": data_x,
+        "yb": data_y,
+        "ab": th_inputs["ab"],  # reuse the previous random value
+    }
+    obj_flattened.update(th_inputs)
+    lin_flattened.linearize()
+
+    # Both objectives should result in the same error and linearizations
+    torch.testing.assert_close(obj_unrolled.error(), obj_flattened.error())
+    torch.testing.assert_close(lin_unrolled.b, lin_flattened.b)
+    torch.testing.assert_close(lin_unrolled.AtA, lin_flattened.AtA)

theseus/core/robust_cost_function.py

@@ -41,9 +41,13 @@
 #       by `weighted_error()` and **NOT** the one returned by
 #       `weighted_jacobians_error()`.
 #
-# Finally, since we apply the weight before the robust loss, we adopt the convention
+# Since we apply the weight before the robust loss, we adopt the convention
 # that `robust_cost_fn.jacobians() == robust_cost_fn.weighted_jacobians_error()`, and
 # `robust_cost_fn.error() == robust_cost_fn.weighed_error()`.
+#
+# The flag `flatten_dims` can be used to apply the loss to each dimension of the error
+# as if it was a separate error term (for example, if one writes a regression problem
+# as a single CostFunction with each dimension being a residual term).
 class RobustCostFunction(CostFunction):
     _EPS = 1e-20
 
@@ -52,6 +56,7 @@ def __init__(
         cost_function: CostFunction,
         loss_cls: Type[RobustLoss],
         log_loss_radius: Variable,
+        flatten_dims: bool = False,
         name: Optional[str] = None,
     ):
         self.cost_function = cost_function
@@ -70,6 +75,7 @@ def __init__(
         self.log_loss_radius = log_loss_radius
         self.register_aux_var("log_loss_radius")
         self.loss = loss_cls()
+        self.flatten_dims = flatten_dims
 
     def error(self) -> torch.Tensor:
         warnings.warn(
@@ -80,9 +86,13 @@ def error(self) -> torch.Tensor:
 
     def weighted_error(self) -> torch.Tensor:
         weighted_error = self.cost_function.weighted_error()
+        if self.flatten_dims:
+            weighted_error = weighted_error.reshape(-1, 1)
         squared_norm = torch.sum(weighted_error**2, dim=1, keepdim=True)
         error_loss = self.loss.evaluate(squared_norm, self.log_loss_radius.tensor)
 
+        if self.flatten_dims:
+            return (error_loss.reshape(-1, self.dim()) + RobustCostFunction._EPS).sqrt()
         # The return value is a hacky way to make it so that
         # ||weighted_error||^2 = error_loss
         # By doing this we avoid having to change the objective's error computation
@@ -107,15 +117,25 @@ def weighted_jacobians_error(self) -> Tuple[List[torch.Tensor], torch.Tensor]:
             weighted_jacobians,
             weighted_error,
         ) = self.cost_function.weighted_jacobians_error()
+        if self.flatten_dims:
+            weighted_error = weighted_error.reshape(-1, 1)
+            for i, wj in enumerate(weighted_jacobians):
+                weighted_jacobians[i] = wj.view(-1, 1, wj.shape[2])
         squared_norm = torch.sum(weighted_error**2, dim=1, keepdim=True)
         rescale = (
             self.loss.linearize(squared_norm, self.log_loss_radius.tensor)
             + RobustCostFunction._EPS
         ).sqrt()
 
-        return [
+        rescaled_jacobians = [
             rescale.view(-1, 1, 1) * jacobian for jacobian in weighted_jacobians
-        ], rescale * weighted_error
+        ]
+        rescaled_error = rescale * weighted_error
+        if self.flatten_dims:
+            return [
+                rj.reshape(-1, self.dim(), rj.shape[2]) for rj in rescaled_jacobians
+            ], rescaled_error.reshape(-1, self.dim())
+        return rescaled_jacobians, rescaled_error
 
     def dim(self) -> int:
         return self.cost_function.dim()
@@ -126,6 +146,7 @@ def _copy_impl(self, new_name: Optional[str] = None) -> "RobustCostFunction":
             type(self.loss),
             self.log_loss_radius.copy(),
             name=new_name,
+            flatten_dims=self.flatten_dims,
         )
 
     @property

theseus/core/vectorizer.py

@@ -13,6 +13,7 @@
 
 from .cost_function import AutoDiffCostFunction, CostFunction
 from .objective import Objective
+from .robust_cost_function import RobustCostFunction
 from .variable import Variable
 
 _CostFunctionSchema = Tuple[str, ...]
@@ -23,6 +24,11 @@ def _fullname(obj) -> str:
         _name = f"{obj.__module__}.{obj.__class__.__name__}"
         if isinstance(obj, AutoDiffCostFunction):
             _name += f"__{id(obj._err_fn)}"
+        if isinstance(obj, RobustCostFunction):
+            _name += (
+                f"__{_fullname(obj.cost_function)}__"
+                f"{_fullname(obj.loss)}__{obj.flatten_dims}"
+            )
         return _name
 
     def _varinfo(var) -> str: