Reg-NF: Efficient Registration of Implicit Surfaces within Neural Fields

Neural fields (NFs) are implicit representations of 3D space that have gained great popularity in recent years [3, 1, 5, 2, 9, 10, 11, 4]. NFs use a small neural network to map any point in a normalised 3D space $\textbf{x}\in\mathbb{R}^{3}$ to the field’s desired output/s (e.g., colour, density, opacity, etc.).

The recent popularity of neural fields is attributed to the introduction of Neural Radiance Fields (NeRFs) [1] which has spawned many radiance field derivatives [3, 12, 11]. The NeRF model is a neural field which maps x and a viewing angle $\textbf{v}\in\mathbb{R}^{2}$ to view-dependant colour $\textbf{c}(\textbf{x},\textbf{v})\in\mathbb{R}^{3}$ and view-independent density $\sigma(\textbf{x})\in\mathbb{R}$. Density represents the likelihood of x hitting anything in space. With this model, any pixel can be represented as a ray (r) described by origin (o) and normalised orientation (v), and the colour of that pixel can be computed through volumetric rendering [1]. NeRF can be trained using only images and associated camera poses by comparing the true colour of a given pixel to the one rendered from the NeRF.

Despite it’s strengths, NeRF’s focus is on visual rendering which can lead to inaccurate geometric representations [9, 5]. Surface neural network models $(S)$  [7, 13, 5, 10, 9], seek to solve this by describing a continuous function $f$ mapping x to the distance to the nearest surface $f(\textbf{x})\in\mathbb{R}$ with the surface located wherever $f(\textbf{x})=0$. These are known to provide smooth and consistent geometric representations [5]. To enable training through purely image data as done by NeRF, some works combine the volumetric representation with implicit surface models [5, 10, 9] and it is these models that will be the focus within our work, outlined in more detail within Section III.

Unlike explicit scene representations, such as voxels, point clouds, meshes, and surfels [14], NFs have attractive properties for robotics as their representations are continuous and memory efficient. Recent work has extended the uses of NFs beyond synthesising novel views to robotics applications. Prior methods [15, 16, 17] demonstrate using the gradient of density from the NFs for collision-free robot motion planning. Several approaches [18, 19, 20] show how continuous representation of NFs can be combined with additional constraints to jointly optimise for grasp and motion planning.

Many algorithms have recently been developed to extend Neural Fields in simultaneous localisation and mapping (SLAM). iMAP[21] is the first method to show NeRF-enabled SLAM with the aid of depth measurements from RGB-D sensors to reconstruct room-size scenes in real-time. NICE-SLAM[22] introduces a hierarchical implicit representation to represent larger scenes, and NICER-SLAM[23] applies neural implicit representations for RGB-only SLAM and shows promising real-time properties.

Registration is the task of estimating the relative transformation between two 3D scene models. It has been extensively studied on explicit representations (e.g., point clouds), whilst registering multiple implicit scene representations remains an underexamined challenge. Nerf2nerf [8] is the first work which demonstrates relative transformation estimation between two NeRF models. However, they rely on human-annotated keypoints for initialisation, which is not practical in robotics applications. DReg-NeRF [24] uses NeRF models then coverts them to an occupancy voxel grid to train 3D CNN followed by attention layers to learn the relations between the pairwise feature grids. DReg-NeRF estimates transformation on explicit and discrete representation. Zero-NeRF[25] performs image to image registration leveraging NeRF representations but in the image space as opposed to between 3D scenes. Moreover, all the exiting prior works assume the scale of the two models are the same. In reality, each model’s coordinate frame will be normalised to fit their specific training conditions. This does not guarantee consistent scale, particularly when considering models trained to represent individual objects and those trained to represent scenes. By comparison, our proposed method does not rely on human-annotated keypoints, and is capable of estimating transformation between two models with arbitrary scales.

In Reg-NF, we provide a technique for aligning the surfaces of two different SDF NFs, by minimising the difference between their surfaces values. Minimisation is optimised for the 6-DoF pose transformation T between the two NFs. We calculate T using a differentiable optimisation function, initialised with an automated procedure. While Reg-NF is generalisable across different use-cases, in this paper we aim to find T between a detected object in a larger scene NF with one in a pre-trained object NF library. We denote $a$ as the notation of an implicit representation of a scene, and $b^{q}$ as the $q$th object-specific implicit model from a database of object neural fields where $q\in\{1,...,Q\}$. Please see Fig. 2 and Fig. 1 for an overview of Reg-NF and our considered use-case pipeline respectively.

Our method assumes we have access to both scene and object NF models. As background information, we provide a brief overview of the NF representation used for our work and its training process.

We utilise volumetric implicit surface fields, specifically NeuS [5] for all NFs in our work. These provide colour $c(\textbf{x},\textbf{v})$ and signed surface distance $f(\textbf{x})$ mapping functions derived from a shared backbone network with separate SDF and colour output heads.

For any given pixel in an image, points at distance ($t$) along the ray emitted from this pixel are represented as $\textbf{r}(t)=\textbf{o}+t\textbf{v}|t\geq 0$. Given a sampling of points along a ray in the range $[t_{1},t_{2},...,t_{n}]$ of increasing magnitude, volumetric rendering can be used to calculate the colour of a given pixel:

where $T_{i}$ is the discrete accumulated transmittance at the $i$th point defined by $T_{i}=\prod_{j=1}^{i-1}(1-\phi_{j})$ and $\phi_{i}$ is the discrete opacity at the $i$th point along the ray. Discrete opacity in NeuS is akin to density in NeRF only it is calculated based on the neural SDF output at the given point along the ray.

Training of the model is done by randomly sampling a batch of ($m$) pixels from a set of training images with known camera poses such that we get training data $D=\{C(\textbf{r}_{k}),\textbf{o}_{k},v_{k}\}$ from which $n$ points are sampled along each ray. Loss is then computed as $L=L_{c}+\lambda L_{e}$ where $L_{c}$ is the colour loss term defined by $L_{c}=\frac{1}{m}\sum_{k}|\hat{C}(\textbf{r}_{k})-C(\textbf{r}_{k})|$, $L_{e}$ is an Eikonal loss regularisation term defined as $L_{e}=\frac{1}{nm}\sum_{k,i}(||f^{\prime}(\textbf{r}_{k}(t_{i}))||_{2}-1)^{2}$, and $\lambda$ is a weighting factor for the Eikonal regularisation. For more details, refer to [5].

We establish neural field registration as an optimisation function that uses gradient decent to find the optimal pose transformation T between two surface fields. The parameters of our optimisation are initialised using $\hat{\textbf{T}}$, and assuming an initial scale factor of $s=1$. Our neural field registration problem can be expressed by the following equation:

where $S_{a}$ and $S_{b}^{q}$ denotes the signed distance representation for implicit models $a$ and $b^{q}$.

In nerf2nerf [8], optimisation is performed over a discrete set of samples collected from a single implicit model, with a robust kernel $\kappa$ used to improve the robustness against outlier samples. The kernel contains learnt parameters $p$ and $\alpha$ which control the decision boundary and impact of outlier samples on registration. In our approach, we consider that the accuracy of optimisation can be improved by collecting samples from both $S_{a}$ and $S_{b}^{q}$ ($A$ and $B^{q}$ respectively), and performing a bidirectional optimisation over these surfaces. Furthermore, we include a regulariser, which is designed to penalise the function when $A$ and $B^{q}$ are deviating from each other.

where $L_{r}$ is our regulariser with a weight factor $w$. Our loss function $L_{s}$ is calculated using the current estimated pose transform T, which we compose by multiplying three transformation matricies together, representing rotation, translation and scale factor components:

where $T$ denotes a 6-DoF translation transformation matrix with components $[t_{x},t_{y},t_{z}]$, $R$ denotes a rotation transformation matrix composed from Euler angles $[r_{r},r_{p},r_{y}]$, and $\sigma$ denotes a scale transformation matrix composed by scale factor $s$. Note that we use a scalar scale factor across all axes to avoid object warping. Our loss function is optimised over the learnt parameters: $(t_{x},t_{y},t_{z},r_{r},r_{p},r_{y},s,p,\alpha)$.

Our registration residuals $r$ are expressed as:

and in the bidirectional case:

Therefore, the residuals are calculating the difference in surface values between SDFs $S_{a}$ and $S_{b}^{q}$. It is important to note that since SDF values are continuous, this loss formulation is differentiable. As SDF values increase in absolute magnitude the further they are away from a surface, the gradient descent will optimise the transformation T to match the surfaces even if the initial samples $A$ and $B^{q}$ have poor initialisation.

Initially, we set samples $A$ and $B^{q}$ equal to the initial surface samples $P_{a}$ and $P_{b}^{q}$. We allow for new samples to be discovered using a variant of the Metropolis-Hastings sampling procedure in nerf2nerf [8]. Our updated sampling algorithm is shown in Algorithm 1. Because SDF values are continuous and not discrete, the sampling strategy needs to be carefully designed to prevent the generation of samples away from true object surfaces.

Our regularisation loss ($L_{r}$) term calculates the mean distance in 3D sample positions between the two different sets of sample points (one from each direction of optimisation); ${A,B^{q}}$. For point correspondence, we assume the nearest 3D point in the other set of sample points ${B^{q}}$ is the ground-truth corresponding point given a particular sample point from ${A}$ (similar to ICP [28]).

where $|A|$ denotes the number of samples in $A$. Given some sample points $A$ from $S_{a}$, and a different set of transformed (by T) sample points $B^{q}$ from $S_{b}^{q}$, after optimisation, their 3D points should be approximately the same. The bidirectional components are balanced by $L_{r}$, which also acts to penalise ‘unmatched points’ between the two models. For example, if one set of samples covers the entirety of a large table in one SDF, and the second set of samples only considers a small proportion of the same table object in the second SDF, then a large disparity will be identified by the regulariser.

We evaluate and compare the performance of Reg-NF and nerf2nerf [8], on our ONR dataset. We also compare against FGR for completeness [29]. Results shown in Table I are the average results across 10 iterations of the experiment to account for randomized factors such as RANSAC. In Table I we see that Reg-NF is typically at least an order of magnitude better than nerf2nerf in terms of $\Delta\textbf{t}$ and is still generally superior in $\Delta\textbf{R}$. We attribute the large increase of errors for nerf2nerf as being primarily due to the inherent scale differences between scene and database object models, for which nerf2nerf has no functionality to handle.

Focusing on Reg-NF, we note that failures can still occur, such as when we match object dc to scene Mix Room or et to scene at-room. In both these cases, we note the cause of failure being poor initialisation that proved inescapable for Reg-NF. A qualitative analysis of the Reg-NF registration can be seen in Fig 4 where coloured versions of NF object library models are substituted into the original scene NF at their poses calculated by Reg-NF.

We conducted an ablation study to investigate the performance of our approach under extreme scale differences between the two neural fields. We use Mix Room for our scene neural field and we attempt to align the fc object within that room to library fc models. For this test, we generated four additional fc models, each resized with scale factors $2,0.5,0.25,0.1$. For these experiments, as convergence takes longer with large scale differences, we ran our optimisation procedure for 1000 iterations. Our results are shown in Table II. We observe that our approach can successfully align the two object fields at a scale factor of 0.5. Interestingly, our learnt scale optimisation can successfully approximate the true scale factor between the scene and library models, even at a scale factor of $2$. In extreme scale scenarios, we observe that the method has room for improvement.

The benefit of multi-view sampling during initialisation is most felt when an object has no distinguishing characteristics or is only partially seen from a single viewing angle taken from the training data. Using a single view this way introduces a high level of variability in object coverage. To show the impact of this, we perform experiments on the ac-room scene, comparing multi-view sampling to using a single challenging view extracted from the training data that observes the object. These views never see the front of chairs and often view only part of the chair, representing worst-case scenarios. The rest of the Reg-NF pipeline is kept consistent and quantitative results are shown in Table III. It is shown that a bad viewpoint with poor samples drastically reduces performance as no meaningful features could be extracted to enable effective registration.

We examine two types of “imperfect scene” to demonstrate practical applications for Reg-NF. The first considers when a robot may not be able to fully traverse a scene to get “full coverage” of an object for the scene’s NF. For this we use the “short” trajectories of our mix-room scene. We can see in Fig. 5 that the scene NF trained from the short trajectory is not able to render the back of the chair clearly as it has no information about that model. Using Reg-NF, we register the object NF for the chair within the scene and can render a clear view of the back of the chair, fully understanding its geometry.

The second application setting considered for Reg-NF is operations with an under-trained scene model. While there are techniques available for speeding up NF training [4], in robotics, hardware limitations may lead to situations where a scene model cannot be fully trained before needing to be used. We therefore investigate the applicability of Reg-NF on highly under-trained scene NFs, trained with 300 steps rather than the 100,000 used in the main experiments. While being an extreme example, we show in Fig. 6 that an accurate alignment and object NF substitution is still feasible for some objects in this setting, enabling a clear render of the object even while the rest of the scene is under-defined.

Finally, we demonstrate the benefits of using a library of pre-trained NF objects for creating new scenes. As all objects in the library are pre-defined and standardised, once Reg-NF derives the transform between a matched object NF and the scene NF, the known relative shapes/poses of objects within the NF library can be used to replace registered scene objects with any other object instance. We demonstrate this within Fig. 7, showing dc being replaced with mc. This opens up potential use-cases in data-driven robotics where, after a scene has been turned into a neural field, objects within that scene can be changed to provide new data based on the layout of the original scene.

The primary limitation of Reg-NF is the initialisation process. This is consistently the most extreme failure case in our work, when an initialisation is so poor that Reg-NF cannot find an optimal solution. This is a problem not faced by nerf2nerf due to the manual keypoint selection. We also acknowledge that current results have assumed a direct mapping from a scene object to a known object model of the same type, an assumption that cannot hold true for all real-world settings. This opens avenues of future work on template warping from a known library of objects.