Update robust loss

theseus/core/robust_cost_function.py

@@ -52,6 +52,7 @@ def __init__(
         cost_function: CostFunction,
         loss_cls: Type[RobustLoss],
         log_loss_radius: Variable,
+        RobustSum: bool = False,
         name: Optional[str] = None,
     ):
         self.cost_function = cost_function
@@ -70,6 +71,7 @@ def __init__(
         self.log_loss_radius = log_loss_radius
         self.register_aux_var("log_loss_radius")
         self.loss = loss_cls()
+        self.RobustSum = RobustSum
 
     def error(self) -> torch.Tensor:
         warnings.warn(
@@ -80,20 +82,33 @@ def error(self) -> torch.Tensor:
 
     def weighted_error(self) -> torch.Tensor:
         weighted_error = self.cost_function.weighted_error()
-        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.RobustSum:
+            squared_norm = torch.sum(weighted_error**2, dim=1, keepdim=True)
+            error_loss = self.loss.evaluate(squared_norm, self.log_loss_radius.tensor)
+            return (
+            torch.ones_like(weighted_error)
+            * (error_loss / self.dim() + RobustCostFunction._EPS).sqrt()
+            )
+        else:
+            squared_norm = weighted_error**2
+            error_loss = self.loss.evaluate(squared_norm, self.log_loss_radius.tensor)
+            return (
+              (error_loss  + RobustCostFunction._EPS).sqrt()
+            )
+        # Inside will compare squared_norm with exp(radius)
+
+        #print(weighted_error, error_loss)
+        #print("scaled", error_loss)
+        #print(self.dim())
         # 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
         # specifically for robust cost functions. The issue for this type of cost
         # function is that the theory requires us to maintain scaled errors/jacobians
         # of dim = robust_fn.cost_function.dim() to do the linearization properly,
         # but the actual error has dim = 1, being the result of loss(||error||^2).
-        return (
-            torch.ones_like(weighted_error)
-            * (error_loss / self.dim() + RobustCostFunction._EPS).sqrt()
-        )
+
 
     def jacobians(self) -> Tuple[List[torch.Tensor], torch.Tensor]:
         warnings.warn(
@@ -107,15 +122,43 @@ def weighted_jacobians_error(self) -> Tuple[List[torch.Tensor], torch.Tensor]:
             weighted_jacobians,
             weighted_error,
         ) = self.cost_function.weighted_jacobians_error()
-        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 [
-            rescale.view(-1, 1, 1) * jacobian for jacobian in weighted_jacobians
-        ], rescale * weighted_error
+        #squared_norm = torch.sum(weighted_error**2, dim=1, keepdim=True)
+        #print(len(weighted_jacobians), weighted_error.shape, weighted_error.dim(), self.cost_function.dim())
+
+        # I do not check the linearization part. I assume they should be correct
+
+        #print(rescale.shape)
+        # The rescale should be reshape as the jocobian shape such that the multipilication works and fulfill the chain rule
+
+
+        if self.RobustSum:
+
+            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 [
+                rescale.view(-1, 1, 1) * jacobian for jacobian in weighted_jacobians
+            ], rescale * weighted_error
+
+        else:
+            squared_norm = weighted_error**2
+            rescale = (
+                self.loss.linearize(squared_norm, self.log_loss_radius.tensor)
+                + RobustCostFunction._EPS
+            ).sqrt()
+
+            for jacobian in weighted_jacobians:
+                rescale_tile = rescale.unsqueeze(2)
+                rescale_tile = torch.tile(rescale_tile, (1, 1, jacobian.shape[2]))
+
+            return [
+                rescale_tile * jacobian for jacobian in weighted_jacobians
+            ], rescale * weighted_error
+
+
 
     def dim(self) -> int:
         return self.cost_function.dim()
@@ -143,4 +186,4 @@ def _supports_masking(self) -> bool:
     @_supports_masking.setter
     def _supports_masking(self, val: bool):
         self.cost_function._supports_masking = val
-        self.__supports_masking__ = val
+        self.__supports_masking__ = val

theseus/core/robust_loss.py

@@ -11,13 +11,14 @@
 
 
 class RobustLoss(abc.ABC):
+    # .suqre() make it easier to compare with residual.square(), which is x
     @classmethod
     def evaluate(cls, x: torch.Tensor, log_radius: torch.Tensor) -> torch.Tensor:
-        return cls._evaluate_impl(x, log_radius.exp())
-
+        return cls._evaluate_impl(x, log_radius.square()) 
+    
     @classmethod
     def linearize(cls, x: torch.Tensor, log_radius: torch.Tensor) -> torch.Tensor:
-        return cls._linearize_impl(x, log_radius.exp())
+        return cls._linearize_impl(x, log_radius.square())
 
     @staticmethod
     @abc.abstractmethod
@@ -49,4 +50,4 @@ def _evaluate_impl(x: torch.Tensor, radius: torch.Tensor) -> torch.Tensor:
 
     @staticmethod
     def _linearize_impl(x: torch.Tensor, radius: torch.Tensor) -> torch.Tensor:
-        return torch.sqrt(radius / torch.max(x, radius) + _LOSS_EPS)
+        return torch.sqrt(radius / torch.max(x, radius) + _LOSS_EPS)