Planning with Learned Dynamic Model for Unsupervised Point Cloud Registration

1 Introduction

Point cloud registration plays an important role in a variety of 3D computer vision applications such as object pose esti-

mation [Wong et al., 2017], Lidar SLAM [Deschaud, 2018], 3D scene reconstruction [Schonberger, 2016]. Given source and target point clouds, it aims to search for an optimal rigid transformation consisting of a rotation matrix and a translation vector so that the transformed source point cloud can align the target point cloud optimally. However, due to the large pose changes and occlusions of objects in real scenarios, point cloud registration is still a challenging problem.

∗Corresponding Author

In past decades, much research efforts have been dedicated to the point cloud registration task. The traditional registration methods such as iterative closest point algorithm (ICP) [Besl and McKay, 1992] and its variants [Chetverikov and Krsek, 2002; Dong et al., 2014] search for the optimal rigid transformation via alternately estimating the correspondences and executing the least squares optimization. Nevertheless, due to the non-convexity of the optimization function, the performance of ICP mainly depends on a good initialized transformation and is prone to the local minima with an improper initialization. To alleviate it, Go-ICP and other variants [Yang et al., 2013; Fitzgibbon, 2003] employ the branch and bound strategy or other heuristic methods to globally search for the transformation. However, these methods usually need much more inference time than the original ICP and just can achieve limited performance improvement. In recent years, benefiting from the discriminative representation ability of deep learning, deep point cloud registration methods [Gojcic et al., 2019; Khoury et al., 2017] can achieve the impressive registration precision. However, their successes mainly rely on large amounts of ground truth corresponding points, which may increase the cost and limit its applications.

Inspired by model-based reinforcement learning [Hafner et al., 2019a], in this paper, we propose a novel unsupervised point cloud registration method by converting the point cloud registration problem into a planning problem. With a learned latent dynamic model, it can efficiently seek the rigid transformation between the source and target point clouds by iteratively narrowing the search space of the transformation via trial and error as shown in Fig. 1. Specifically, the constructed dynamic model contains a transformation network and evaluation network. The transformation network captures the feature changes when a point cloud is transformed, and is used to predict the new transformed feature of the point cloud after performing a rigid transformation on it. The evaluation network captures the geometric difference between the transformed source point cloud and target point cloud in the feature space, which can predict the alignment precision as the reward in planning. We formulate a transformation consistency loss on the transformation network and a geometry consistency loss on the evaluation network to train the dynamic model. Then, based on the learned model, we employ the cross-entropy method (CEM) [Botev et al., 2013] to globally search for the optimal rigid transformation. CEM is a sampling-based planning algorithm, which alternately samples multiple transformations from a distribution and selects some "elite" ones with the high evaluation scores to update the sampling distribution. The sampled transformations in each iteration are expected to approach to the optimal solution progressively and the mean of the elite transformations in the last iteration is viewed as the final solution, i.e., the transformation between the source and target point clouds. Experimental results on benchmark datasets demonstrate that the proposed registration algorithm can yield good performance.

In summary, the contributions of this paper are as follows:

• We model the point cloud registration task as a Markov decision process and propose a latent dynamic model of point clouds for planning.
• We perform planning with the learned latent dynamic model for unsupervised point cloud registration.
• Extensive experiments on the standard benchmark datasets verify the outstanding registration performance of the proposed method.

2 RelatedWork

Traditional point cloud registration algorithms. Iterative closest point (ICP) algorithm [Besl and McKay, 1992] is a popular method for fine registration, which alternates between estimating the correspondences and updating the transformation based on the least squares optimization. However, with an improper initial transformation, it easily falls into the local minima. To relieve it, Go-ICP [Yang et al., 2013] is proposed to search for the optimal transformation by combining ICP with the branch-and-bound scheme. In addition, RANSAC based methods [Fischler, 1981; Chum and Matas, 2005] are also widely-used for coarse registration where wide bases are sampled for aligning the point cloud pair. 4PCS [Aiger et al., 2008] is a representative RANSAC based method, which proposes to sample four points with an approximately co-planar layout as a base set and identify the corresponding sets in the point cloud pair by comparing the intersectional diagonal ratios of the base set. Moreover, other RANSAC variant [Mellado et al., 2014] also shows good performance.

Learning based point cloud registration algorithms. In recent years, more and more attention has been paid to deep learning based point cloud registration. Benefitting from that PointNet [Qi et al., 2017] can directly extract robust feature of unordered point cloud with deep neural network, this research direction has obtained great progress. PRNet [Wang, 2019] is an iterative method that utilizes the keypoint detection to handle the partially overlapping point cloud pair in a self-supervised way. Instead, DCP [Wang and Solomon, 2019] performs the SVD module to estimate the transformation based on a soft matching between the point cloud pair. RPM-Net [Yew and Lee, 2020] also realizes the registration based on the matching map generated from the differentiable Sinkhorn layer and annealing. Yang et al. [Yang et al., 2020] proposed an unsupervised learning method for finding the correspondences based on the cycle consistency. In addition, other learning-based methods including [Aoki et al., 2019; Li et al., 2020] also exhibit impressive performance.

Planning based optimal decision. Planning based decision algorithm is important and well-studied in the optimal control field. Typically, those methods search for the optimal action via repeatedly performing forward simulation and future state evaluation. For discrete control task, Monte-Carlo tree search (MCTS) [Browne et al., 2012] is one of the most widely used method whose each simulation procedure contains selection, evaluation, expansion and backing-propagation. To handle exploration-exploitation dilemma in MCTS, UCT and its variants [Coulom, 2006; Rosin, 2011] add an additional exploration bonus on original estimated reward for each action. For the continuous action task, the cross entropy method (CEM) [Botev et al., 2013] is a powerful planning method, which can obtain the optimal action via gradually narrowing the searching space in a trialand-error way. To improve the robustness of CEM, Liu et al. [Liu et al., 2020] proposed to perform CEM with model uncertainty and constraints. Furthermore, the differentiable CEM [Amos and Yarats, 2020] is proposed to allow the endto-end training. Other variants or applications [Moss, 2020; Hafner et al., 2019b] also show its great potential.

3 Approach

In this section, we present the proposed point cloud registration method by planning with the latent dynamic model. Fig. 2 illustrates the proposed framework, which consists of two components: learning latent dynamic model of point clouds and planning for point cloud registration. In Section 3.1, we define the point cloud registration problem with the Markov decision process. In Sections 3.2 and 3.3, we present the details of learning the latent dynamic model and planning for point cloud registration, respectively.

3.1 Problem Setting

Given the source point cloud X = {x1, x2, . . . xN} and target point cloud Y = {y1, y2, . . . yN}, where xi ∈ R3 and yi ∈ R3 , i = 1, 2, . . . ,N, the point cloud registration task aims to find the rigid transformation between X and Y, containing the rotation R ∈ SO(3) and the translation t ∈ R3 . The rigid transformation in 3D Euclidean space is defined as rd `N

Action space. We define the rigid transformation as the action. The action space A is the rigid transformation space of point clouds SE(3) (SO(3) × R3 ). To reduce the dimension of the search space, we use the euler-angle representation e = [e1, e2, e3] ∈ [-π, π]3 to encode the rotation R ∈ SO(3). Therefore, the action a ∈ A is represented by [e, t] and R(e) is the corresponding rigid rotation matrix of euler-angle representation e.

Latent transformation function. In the 3D Euclidean space, after performing the action a, the point cloud X will be transformed to X′ via the rigid transformation. Similarly, in the latent state space, the latent transformation function f z θ : Z ×A → Z aims to infer the next latent state ˆzX′ by performing action a at the state zX. It is defined as below:

where [·, ·] denotes the concatenation operation. Since the transformation function in the registration task is nonlinear, we utilize a deep network with parameters θ to model it.

Latent reward function. For the source point cloudX and target point cloud Y, given their latent states zX and zY , we use the output of an evaluation network as the reward signal after executing the action a at the latent state zX. Specifically, the latent reward function r : Z ×A×Z → R is defined as:

where the evaluation network dψ : Z ×Z → R estimates the alignment precision between X′ and Y in the feature space. The point cloud registration task aims to find the optimal policy a∗ that can maximize the reward:

3.2 Learning Latent Dynamic Model

Our dynamic model of point clouds is constructed in the feature space of point clouds, where we employ an encoderdecoder network to extract deep features (latent states) of point clouds so that the local geometric structures of point clouds can be characterized well. As shown in Fig. 2 (a), by fusing the feature of the source point cloud and the executed action, we first form a transformation network to predict the latent state of the transformed source point cloud. We then employ an evaluation network to evaluate the alignment precision between the transformed source point cloud and the target point cloud in the feature space. The evaluated alignment precision between point clouds is used as the reward in planning for point cloud registration. By performing a set of actions in trial and error to maximize the reward during the registration process, we can seek the action with the highest reward, i.e., the transformation between a pair of point clouds (please refer to Section 3.3).

Before effectively training our dynamic model, we need to generate training samples at first. For each pair of point clouds 〈X,Y〉, we sample multiple actions {a = [e, t]}, where each action sample a follows a normal distribution N (0, σ2 I) over the action space A. Next, we perform each sampled action on X and obtain the transformed point clouds {X′ } where X′ = f 3d (X, R(e), t). Thus, we can obtain the generated training samples S = {X, a, X′ , Y}.

During the model learning, we jointly train an auto-encoder and a latent dynamic model. The auto-encoder is expected to learn the discriminative feature representation (latent state) of point clouds. The encoder ϕenc φ and decoder ϕdec ω modules in the auto-encoder are trained via the reconstruction loss:

where L(X) = Dcd (ϕdec ω (ϕenc φ (X)) ,X). Dcd is the Chamfer distance between X and Y, Dcd (X, Y) =

Notably, once the encoder and decoder are trained, the encoder is only used for testing.

For the source point clouds X , the transformation network f z θ is expected to capture the feature changes when a point cloud is transformed by performing the action on them. That is the predicted transformed state f z θ (zX, a) is as consistent as possible to the state zX′ of the transformed point clouds. Therefore, we formulate a transformation consistency loss Ltrans to train our transformation network:

In our dynamic model, the evaluation network dψ aims to predict the alignment precision between a pair of point clouds with their states, which can characterize the geometric differences of point clouds in the feature space. We encourage the geometric differences of point clouds in the feature and spatial spaces to be consistent. To this end, we propose a geometric consistency loss Leval to train our network:

where I (X,Y) = (dψ (zX, zY ) -Dcd (X,Y))2 , dψ (zX, zY ) is the output of the evaluation network. Notably, we utilize the Chamfer distance to quantify the alignment precision between a point clouds pair as in [Yew and Lee, 2020], since as a distance metric, it tends to obtain the lower value when the point clouds are better aligned.

Finally, we jointly train the encoder-decoder and latent dynamic model via the following loss function:

Algorithm 1 CEM based Registration with Learned Latent Dynamic Model

Input: Encoder ϕenc φ , latent transformation network f z θ , latent evaluation network dψ, iterations T , sampled transformations N, top transformations K, initial Gaussian distribution N (µ0, σ2 0 I), source and target point clouds X and Y.

Output: Estimated rigid transformation between X and Y.

end for

Sort the rewards {ri} and select top transformations K with K largest rewards.

Update the sampling distribution with K via Eq. 9.

end for

Return: µT

3.3 Planning for Point Cloud Registration

As shown in Fig. 2 (b), once the dynamic model of point clouds is trained, given the source and target point clouds, by performing actions on the source point clouds, we can maximize the rewards (i.e., the outputs of the evaluation network) to obtain the optimal action for registration. In this subsection, we employ the cross entropy method (CEM) to obtain the desired transformation for registration. CEM is an efficient sampling-based planning algorithm that searches for the optimal solution via iteratively narrowing the searching region in a trial-and-error way. As outlined in Algorithm 1, in the t-th iteration (0 ≤ t < T - 1), it contains three steps:

• (a) Random transformation sampling. We randomly sample N rigid transformations {a(1) t , . . . , a(N) t } from a parameterized Gaussian distribution N (µt, σ2 t I) over the action space A. Since the dimension of the action space is 6, the mean µt and variance σ2 t are represented by (µt,1, . . . , µt,6) and (σ2 t,1 , . . . , σ2 t,6 ), respectively.
• (b) Top-K transformation selection. From all sampled N rigid transformations, with the trained transformation network and evaluation network, we choose top-K ones with the highest rewards, denoted as K = {a(i) t | r (zX, a(i) t , zY ) = θ ([zX, a]) , zY ]) ≥ βt,1 ≤ i ≤ N}, where βt is the

• (c) Distribution updating. This step aims to update the parameters of the current sampling distribution so that in the next iteration we can sample "better" transformations from the updated distribution N (µt+1, σ2 t+1 I):

where j = 1, . . . , 6. In fact, Eq. 9 is the closedform solution of the maximum-likelihood estimation problem over the top-K transformations, that is µt+1, σt+1 =

argmaxµ,σ ∑a∈K g (a; µ, σ), where g (·; µ, σ) is the probability density function of N (µ, σ2 ). By re-fitting the sampling distribution to the top-K samples, the updated sampling distribution tends to focus on more promising transformation region and the sampled transformations in the next iteration are expected to obtain the higher rewards. In the last iteration, we utilize the mean µT of the distribution to estimate the optimal solution.

4 Experiments

4.1 Implementation Details

For feature extraction of point clouds, we employ the DGCNN [Wang et al., 2019] as encoder and the decoder of FoldingNet [Yang et al., 2018] as decoder in our autoencoder. The latent transformation and evaluation networks are modeled by two 3-layer MLPs and more details are described in Appendix 2. We use the output of the encoder as the latent state of the point cloud whose dimension is set to 512. For CEM, the numbers of iterations T , candidates N and elite candidates K in CEM are set to 10, 1000 and 25, respectively. The σ used in training samples generation in Section 3.2 is set to 1. Our method is implemented in PyTorch and the model is trained with the Adam optimizer , where the learning rate and weight decay are set to 0.0001 and 0.0005. For simplicity, we denote our CEM based latent planning method for registration as LatentCEM.

4.2 Evaluation onModelNet40

We perform performance comparison with four traditional registration methods including ICP [Besl and McKay, 1992], Go-ICP [Yang et al., 2013] , Super4PCS [Mellado et al., 2014] and FGR [Zhou et al., 2016], and three deep learning methods including fully-supervised PointNetLK [Aoki et al., 2019], fully-supervised IDAM [Li et al., 2019] and unsupervised FMR [Huang et al., 2020] on the ModelNet40 dataset. The parameter settings of the traditional methods can be seen in Appendix 3. The ModelNet40 dataset contains 12, 311 mesh CAD models from 40 object categories. We select 2, 468 models for testing and uniformly sample 784 points from the outer surface of each model, where selected points in each model are centered and rescaled to fit within the unit sphere. We utilize the root mean square error (RMSE) and mean absolute error (MAE) between ground truth and predicted transformations measure the quality of registration. All angular error results are converted to the units of degrees. The target point clouds are obtained by transforming the source point clouds via a randomly generated transformation matrix. The rotation and translation along each axis uniformly lie in the [0, 45◦ ] and [-0.5, 0.5]. Notably, no ground-truth transformation is involved in our model training and thus our method is unsupervised. Fig. 3 shows the qualitative comparison and more results can be seen in Appendix 1.

Train and test on the same categories. In this experiments, the samples for training our model and other three deep learning methods are from the same categories as the testing samples. Notably, and the groud-truth transformations are required for PointNetLK and IDAM. As shown in Table 1, benefitting from the learned precise dynamic model and the effective CEM based search policy, our method can obtain the highest precision in terms of the translation error while the second highest precision in terms of the rotation error. And, although the RMSE of our method in rotation is worse than that of IDAM, the IDAM is fully-supervised, which needs amounts of transformation labels as the supervision signal while our method is unsupervised and thus can largely reduce the training cost.

Gaussian noise. To verify the robustness of our proposed point cloud registration method, we train our model on the noise-free data of ModelNet40 and test it on a new test dataset disturbed by Gaussian noise. Following [Wang, 2019], the noisy data is randomly sampled from a Gaussian distribution with mean 0 and standard deviation 0.01, and then clipped to [-0.05, 0.05]. Table 1 shows that our method is robust to the noise and can obtain superior score than other methods.

4.3 Evaluation on 7Scene

We further evaluate our method on the 7Scene dataset, which is a real indoor dataset containing seven scenes: Chess, Fires, Heads, Office, Pumpkin, Redkitchen and Stairs. Following [Huang et al., 2020], We project the depth images into point clouds, and generate 296 scans for training and 57 scans for testing through the similar data processing in ModelNet40 dataset. Here, the target point cloud is also generated via performing a randomly generated transformation on the source point cloud. Table 1 summarizes the comparison results with other methods on the 7Scene dataset and our method can still yield good performance compared to other methods.

4.4 Ablation Study and Analysis

Encoder and decoder setting. For our auto-encoding framework, we conduct ablation studies on ModelNet40 with Computational efficiency. We also compare the inference time of different methods on the ModelNet40 dataset. The running time is measured in millisecond with a batch size of 1, averaged over the entire test set on a desktop computer with an Intel I5-8400 CPU and Geforce RTX 2080Ti GPU. Notably, the traditional methods are executed on CPU while deep learning methods on GPU. The inference time of our method is 75.638 and the inference time of other methods are 8.168 (ICP), 2739.588 (Go-ICP), 6128.691 (Super4PCS), 52.203 (FGR), 386.244 (FMR), 180.645 (PointNetLK) and 41.401 (IDAM), respectively. Fig. 4 shows the running time

... In this paper, we proposed a novel registration method by planning with a learned latent dynamic model. In detail, we first designed an effective training framework to learn the latent dynamic model consisting of a transformation function and a reward function. Then, with the learned dynamic model, we utilized the trial-and-error based cross entropy method to search for the optimal transformation, where the dynamic model is utilized to evaluate the output of each trial and guide the evolution of CEM. Extensive experimental results on ModelNet40 and 7Scene benchmark datasets demonstrate that our method can yield good performance.

Acknowledgments

This work was supported by the National Science Fund of China (Grant Nos. U1713208, 61876084, 61876083), Pro-

of our method in the cases of different iterations. From this

Inputs t=1

t=2

t=3

t=4

t=5

t=6 t=7 t=10 ... gram for Changjiang Scholars.