Rigidity embedding computed in RAFT3D

[justachetan]

Hi @zachteed,

I am interested in visualizing the rigidity embedding computed during the forward pass through the network. It is denoted by $\mathbf{V}$ for the vanilla RAFT3D model and optimized using Eq. 17 for the bi-Laplacian version in the paper. My current hunch is that the rigidity embedding is named ae in the code (in raft3d_bilaplacian.py). Could you kindly confirm if that is indeed correct? I am confused because the paper says that $\mathbf{V} \in \mathbb{R}^{H \times W \times C}$ but in the code, ae is initialized as a $\frac{H}{8} \times \frac{W}{8} \times 16$ matrix.

Also, I could not understand from the paper how the rigidity embedding exactly affects the output SE(3) field. Could you please explain how it factors into the final prediction or is it just predicted as an auxiliary output for regularization purposes?