Inaccurate Results

[btilmon]

Hello, I am doing a very simple PyTorch test with the global alignment. I have two point clouds with significant overlap from a video and I would like to regress the relative pose between them with roma like dust3r did. The gradients seem incorrect on my data though and I cannot get an accurate alignment after several modifications. I am able to get a good alignment using Procrustes SVD method. The point clouds are unnormalized in metric units from a depth sensor. Do you perhaps have any recommendations on this code? I have confirmed learning rate and number of iterations are not a factor.

# taken at time i
x = point_cloud[0,0].view(-1,3)

# taken at time i+1
y = point_cloud[0,1].view(-1,3)

rel = torch.eye(4, device=self.device)
rot = rel[:-1, :-1]
trans = rel[:-1, -1]

rot = torch.nn.Parameter(rot, requires_grad=True)
trans = torch.nn.Parameter(trans, requires_grad=True)

optimizer = torch.optim.Adam([rot, trans], lr=1e-3)
progress_bar = tqdm(range(100)) 

x = torch.cat([x, torch.ones_like(x[..., :1])], dim=-1)
y = torch.cat([y, torch.ones_like(y[..., :1])], dim=-1)

for step in progress_bar:
    optimizer.zero_grad()
    Q = roma.rotmat_to_unitquat(rot)
    rel_roma_rigid = roma.RigidUnitQuat(Q, trans).normalize().to_homogeneous()
    new = x @ rel_roma_rigid
    loss = F.mse_loss(new, y)
    loss.backward()
    optimizer.step()

[btilmon]

Got it working by getting correspondences first and then following this approach:

rotvec = torch.nn.Parameters(torch.zeros(3), requires_grad=True)
...
# In the optimization loop, keeping x and y 3D point clouds
R = roma.rotvec_to_rotmat(rotvec)
T = roma.Rigid(R, trans)
# Transformed point cloud
new = T.apply(x)