From fa9f1a79a5a5e31ecf373007c602eb57834ae4ad Mon Sep 17 00:00:00 2001 From: Zitong Zhan Date: Sun, 18 Oct 2026 22:57:52 -0400 Subject: [PATCH 1/2] Fix SE3 translation gradients in PyPose ambient mode --- bae/utils/pypose_ambient_grad.py | 8 ++++ tests/autograd/test_pypose_ambient_grad.py | 45 ++++++++++++++++++++++ tests/autograd/test_pypose_compile.py | 2 +- 3 files changed, 54 insertions(+), 1 deletion(-) diff --git a/bae/utils/pypose_ambient_grad.py b/bae/utils/pypose_ambient_grad.py index 1ff3642..8aa5de6 100644 --- a/bae/utils/pypose_ambient_grad.py +++ b/bae/utils/pypose_ambient_grad.py @@ -281,6 +281,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 +298,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 diff --git a/tests/autograd/test_pypose_ambient_grad.py b/tests/autograd/test_pypose_ambient_grad.py index a50a7d0..5564555 100644 --- a/tests/autograd/test_pypose_ambient_grad.py +++ b/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 = [ @@ -96,3 +97,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) diff --git a/tests/autograd/test_pypose_compile.py b/tests/autograd/test_pypose_compile.py index a2933c4..9c20add 100644 --- a/tests/autograd/test_pypose_compile.py +++ b/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( From fa1ecfc8327178dd80b51fdd760b11645bcb20e7 Mon Sep 17 00:00:00 2001 From: Zitong Zhan Date: Mon, 19 Oct 2026 01:50:32 -0400 Subject: [PATCH 2/2] Fix ambient Log gradients at identity and half turns --- bae/utils/pypose_ambient_grad.py | 22 +++++--- tests/autograd/test_pypose_ambient_grad.py | 64 ++++++++++++++++++++++ 2 files changed, 78 insertions(+), 8 deletions(-) diff --git a/bae/utils/pypose_ambient_grad.py b/bae/utils/pypose_ambient_grad.py index 8aa5de6..ce25ab8 100644 --- a/bae/utils/pypose_ambient_grad.py +++ b/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 diff --git a/tests/autograd/test_pypose_ambient_grad.py b/tests/autograd/test_pypose_ambient_grad.py index 5564555..37b74e5 100644 --- a/tests/autograd/test_pypose_ambient_grad.py +++ b/tests/autograd/test_pypose_ambient_grad.py @@ -73,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")