Fix SE3 translation and Log gradients in PyPose ambient mode

bae/utils/pypose_ambient_grad.py

@@ -82,14 +82,20 @@ def _so3_log_forward(input: torch.Tensor) -> torch.Tensor:
     w_abs = torch.abs(w)
     v_larger_than_eps = v_norm > eps
     w_larger_than_eps = w_abs > eps
-    idx1 = v_larger_than_eps & w_larger_than_eps
-    idx2 = v_larger_than_eps & (~w_larger_than_eps)
-    idx3 = ~v_larger_than_eps
-
-    factor = torch.zeros_like(v_norm, requires_grad=False)
-    factor = factor + idx1 * torch.nan_to_num(2.0 * torch.atan(v_norm / w) / v_norm)
-    factor = factor + idx2 * torch.nan_to_num(_pm(w) * torch.pi / v_norm)
-    factor = factor + idx3 * torch.nan_to_num(2.0 * (1.0 / w - v_norm * v_norm / (3 * w**3)))
+    # Autograd differentiates inactive branches too: masking or nan_to_num
+    # after a division by zero still lets NaNs enter the backward pass.
+    # Protect denominators before evaluating any branch (identity: v=0;
+    # half turn: w=0), preserving the existing principal-Log convention.
+    safe_v_norm = torch.where(v_larger_than_eps, v_norm, torch.ones_like(v_norm))
+    safe_w = torch.where(w_larger_than_eps, w, torch.ones_like(w))
+    regular = 2.0 * torch.atan(v_norm / safe_w) / safe_v_norm
+    half_turn = _pm(w) * torch.pi / safe_v_norm
+    small = 2.0 * (1.0 / safe_w - v_norm * v_norm / (3 * safe_w**3))
+    factor = torch.where(
+        v_larger_than_eps,
+        torch.where(w_larger_than_eps, regular, half_turn),
+        small,
+    )
     return factor * v
 
 
@@ -281,6 +287,13 @@ def _ambient_se3_inv(self, X):
         X = X.tensor() if isinstance(X, lt.LieTensor) else X
         return lt.LieTensor(_se3_inv_forward(X), ltype=lt.SE3_type)
 
+    def _ambient_se3_translation(self, X):
+        X = X.tensor() if isinstance(X, lt.LieTensor) else X
+        # PyPose's custom translation backward returns a padded tangent gradient.
+        # Keep ambient derivatives here; BAE applies the retraction Jacobian once
+        # at the parameter boundary, after differentiating the full residual.
+        return X[..., :3]
+
     lt.SO3Type.Log = _ambient_so3_log
     lt.SE3Type.Log = _ambient_se3_log
     lt.so3Type.Exp = _ambient_so3_exp
@@ -291,6 +304,7 @@ def _ambient_se3_inv(self, X):
     lt.SE3Type.Mul = _ambient_se3_mul
     lt.SO3Type.Inv = _ambient_so3_inv
     lt.SE3Type.Inv = _ambient_se3_inv
+    lt.SE3Type.translation = _ambient_se3_translation
 
     _PATCH_INSTALLED = True
     return True

tests/autograd/test_pypose_ambient_grad.py

@@ -12,6 +12,7 @@
     sys.path.insert(0, str(_REPO_ROOT))
 
 from bae.utils.pypose_ambient_grad import install_pypose_ambient_grad_monkeypatch  # noqa: E402
+from bae.utils.parameter import trim_parameter_jacobian_values  # noqa: E402
 
 
 pytestmark = [
@@ -72,6 +73,70 @@ def test_pypose_se3_log_monkeypatch_produces_ambient_gradient():
     torch.testing.assert_close(jac, jac_fd, rtol=1e-7, atol=1e-7)
 
 
+@pytest.mark.parametrize("group", [pp.SO3, pp.SE3])
+def test_ambient_log_identity_and_mixed_batch_gradients(group):
+    install_pypose_ambient_grad_monkeypatch()
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    quaternions = torch.tensor(
+        [[0., 0., 0., 1.], [0., 0., 0., -1.],
+         [0., 0., 0., 1.0000011496537984],
+         [1e-8, -2e-8, 3e-8, 1.], [0.2, -0.3, 0.1, 0.9]],
+        device=device, dtype=torch.float64,
+    )
+    values = quaternions
+    if group is pp.SE3:
+        translation = quaternions.new_tensor([0.3, -0.2, 0.5]).expand(5, 3)
+        values = torch.cat((translation, quaternions), dim=-1)
+
+    def log(x):
+        return group(x).Log().tensor()
+
+    # vmap/jacrev is the path used to differentiate individual sparse factors.
+    jac = torch.func.vmap(torch.func.jacrev(log))(values)
+    assert torch.isfinite(jac).all()
+    for value, derivative in zip(values, jac):
+        torch.testing.assert_close(derivative, _fd_jacobian(log, value), rtol=1e-6, atol=1e-6)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
+def test_ambient_so3_log_half_turn_is_finite(dtype):
+    install_pypose_ambient_grad_monkeypatch()
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    quaternion = torch.tensor([1., 0., 0., 0.], device=device, dtype=dtype, requires_grad=True)
+    result = pp.SO3(quaternion).Log().tensor()
+    torch.testing.assert_close(result, quaternion.new_tensor([torch.pi, 0., 0.]))
+    # The principal Log is discontinuous here. Check numerical finiteness of
+    # the selected branch, not agreement with a derivative across the cut.
+    result.sum().backward()
+    assert torch.isfinite(quaternion.grad).all()
+
+
+def test_ambient_pgo_log_at_identity_has_correct_local_jacobian(monkeypatch):
+    monkeypatch.setenv("BAE_USE_PYPOSE_AMBIENT_GRAD", "1")
+    install_pypose_ambient_grad_monkeypatch()
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    values = torch.tensor(
+        [[0., 0., 0., 0., 0., 0., 1.], [0.3, -0.2, 0.5, 0., 0., 0., 1.]],
+        device=device, dtype=torch.float64,
+    )
+    nodes = pp.SE3(values)
+    measurement = nodes[0].Inv() @ nodes[1]
+
+    def residual(x):
+        poses = pp.SE3(x)
+        return (measurement.Inv() @ poses[0].Inv() @ poses[1]).Log().tensor()
+
+    torch.testing.assert_close(residual(values), values.new_zeros(6))
+    ambient = torch.func.jacrev(residual)(values)
+    local = trim_parameter_jacobian_values(nodes, ambient.transpose(0, 1)).transpose(0, 1)
+
+    def perturbed_residual(delta):
+        return residual((pp.se3(delta).Exp() @ nodes).tensor())
+
+    local_fd = _fd_jacobian(perturbed_residual, values.new_zeros(2, 6))
+    torch.testing.assert_close(local, local_fd, rtol=1e-7, atol=1e-7)
+
+
 def test_pypose_se3_inv_monkeypatch_produces_ambient_gradient():
     install_pypose_ambient_grad_monkeypatch()
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -96,3 +161,47 @@ def test_pypose_se3_exp_monkeypatch_produces_ambient_gradient():
     jac_fd = _fd_jacobian(lambda x: pp.se3(x).Exp().tensor(), algebra)
 
     torch.testing.assert_close(jac, jac_fd, rtol=1e-7, atol=1e-7)
+
+
+@pytest.mark.parametrize("batch_shape", [(), (2,), (2, 3)])
+def test_pypose_se3_translation_monkeypatch_produces_ambient_gradient(batch_shape):
+    install_pypose_ambient_grad_monkeypatch()
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    pose = torch.tensor(
+        [0.3, -0.2, 0.5, 0.0, 0.0, 0.0, 1.0], device=device, dtype=torch.float64
+    ).expand(*batch_shape, 7).clone()
+
+    def translation(x):
+        return pp.SE3(x).translation()
+
+    torch.testing.assert_close(translation(pose), pose[..., :3])
+    jac = torch.func.jacrev(translation)(pose)
+    jac_fd = _fd_jacobian(translation, pose)
+    torch.testing.assert_close(jac, jac_fd, rtol=1e-7, atol=1e-7)
+
+
+def test_se3_relative_translation_localized_jacobian_matches_pose_perturbations(monkeypatch):
+    monkeypatch.setenv("BAE_USE_PYPOSE_AMBIENT_GRAD", "1")
+    install_pypose_ambient_grad_monkeypatch()
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    algebra = torch.tensor(
+        [[0.3, -0.2, 0.5, 0.2, 0.3, -0.1], [-0.4, 0.8, 0.2, -0.3, 0.1, 0.4]],
+        device=device, dtype=torch.float64,
+    )
+    nodes = pp.se3(algebra).Exp().detach()
+
+    def relative_translation(x):
+        poses = pp.SE3(x)
+        return (poses[0].Inv() @ poses[1]).translation()
+
+    # BAE differentiates stored coordinates, then converts each parameter block
+    # once to the six-dimensional left perturbation used by the optimizer.
+    ambient = torch.func.jacrev(relative_translation)(nodes.tensor())
+    local = trim_parameter_jacobian_values(nodes, ambient.transpose(0, 1)).transpose(0, 1)
+
+    def perturbed_translation(delta):
+        perturbed = pp.se3(delta).Exp() @ nodes
+        return relative_translation(perturbed.tensor())
+
+    local_fd = _fd_jacobian(perturbed_translation, torch.zeros_like(algebra))
+    torch.testing.assert_close(local, local_fd, rtol=1e-7, atol=1e-7)

tests/autograd/test_pypose_compile.py

@@ -100,7 +100,7 @@ def residual(a, b, p):
         pose_a = pp.SE3(a)
         pose_b = pp.SE3(b)
         relative = (pose_a.Inv() @ pose_b).Log().tensor()
-        return torch.cat((relative, pose_a.Act(p)), dim=-1)
+        return torch.cat((relative, pose_a.Act(p), pose_a.translation()), dim=-1)
 
     eager = residual(lhs, rhs, points)
     eager_gradients = torch.autograd.grad(