RLSAC: Reinforcement Learning enhanced Sample Consensus
for End-to-End Robust Estimation

Robust estimation can be considered as generating a good hypothesis $h$ given a set of data $\chi=\left\{{{{\rm{x}}_{i}}}\right\}_{i=1}^{N}$, which may be disturbed by noise. For instance, the hypothesis $h$ could represent the parameters of a 2D line, and $\chi$ could be all the points in the 2D plane. Or $h$ could be the fundamental matrix that represents the epipolar geometry of a pair of images, and $\chi$ could be all the correspondences. Similarly, the data contained in ${{\rm{x}}_{i}}$ varies with the task. As a commonly used robust estimation method, the sampling consensus method samples $n$ minimum sets $\mathcal{M}$ from $\chi$. The size $m$ of a minimum set varies depending on the task. For example, when fitting a 2D line, the size $m=2$. Then, these minimum sets $\mathcal{M}$ are solved by the minimum solver $S$ to output hypotheses:

$$H=\left\{{S\left({{\mathcal{M}_{j}}}\right)\left|{{{\mathcal{M}_{j}}}\in M,j=1,2,...,n}\right.}\right\}.$$

(1)

Next, the hypotheses $H$ are evaluated using a scoring function $f$. The residuals of all data points $\chi$ can be computed and used to calculate the inlier ratio, which is commonly used as a metric for the hypothesis quality [16]. Finally, the hypothesis with the highest score is chosen as the best hypothesis $h_{Best}$:

$${h_{Best}}=\mathop{\arg\max}\limits_{h\in H}f\left({h,\chi}\right).$$

(2)

By solving the problem using several minimum sets $\mathcal{M}$, the model can achieve robust estimation, which is less sensitive to outliers and can lead to more accurate results.

The RANSAC [16] algorithm performs random sampling to sample the minimum sets $\mathcal{M}$, which cannot fully exploit the data features. As shown in Figure 1, One alternative is to consider the sampling of a minimum set as an action taken by an agent based on the current state, within the reinforcement learning framework:

$$a_{t+1}\sim\pi_{\phi}\left(a_{t}\mid s_{t}\right).$$

(3)

Here, the policy network $\pi_{\phi}$ depends on the weights $\phi$. The action $a_{t+1}$ is sampled by the policy $\pi_{\phi}$ from the state $s_{t}$ and $a_{t}$ at time $t+1$, which can also be viewed as the process of sampling a minimum set $\mathcal{M}_{j}$ from all data $\chi$. So the Eq. 2 can be rewritten that the best hypothesis is selected by a reinforcement learning enhanced sample consensus method:

$$h_{Best}=\mathop{\arg\max}\limits_{t=1,2,...,j}f\left({S\left({{a_{t}}}\right),\chi}\right).$$

(4)

In addition, the policy $\pi_{\phi}$ is a trainable neural network, and the state $s_{t}$ contains the features of the data $\chi$. Therefore, the minimum set can be selected based on sufficient learned knowledge of the data features, rather than being chosen randomly. Furthermore, the evaluation result $f\left({S\left({{a_{t}}}\right),\chi}\right)$ of the minimum solver $S$ can serve as the reward in reinforcement learning to achieve unsupervised learning.

With the problem formulation of modeling the process of sampling consensus as a reinforcement learning problem, we introduce RLSAC, which is illustrated in Figure 2. Although the fundamental matrix estimation task is used as an example, the pipeline is also applicable to other robust estimation tasks. Firstly, RLSAC initiates multiple episodes for the correspondences of a pair of images $\chi$ at \raisebox{-0.7pt}{1}⃝. An episode contains many steps. At the start of each episode, a random sampling of the minimum set $\mathcal{M}_{0}$ is performed to generate the initial state ${s_{0}}$ through state transition at \raisebox{-0.7pt}{2}⃝. Specifically, the initial state ${s_{0}}$ is obtained by concatenating each data point with the memory features, which are consisting of action, residual, and historical features (see Section 3.5).

Next, the initial state ${s_{0}}$ is input into the sampling consensus loop. The agent receives the current state ${s_{t}}$ and feeds it into the policy network $\pi_{\phi}$ at \raisebox{-0.7pt}{3}⃝, as shown in Section 3.3. The network generates a probability for each data point. The $m$ points with the highest probabilities are selected as the minimum set $\mathcal{M}_{t}$, instead of random sampling.

Then, the minimum set $\mathcal{M}_{t}$ serves as the action $a_{t}$ for the agent, which generates a hypothesis ${h_{t}}$ in the environment at \raisebox{-0.7pt}{4}⃝, as shown in Section 3.4. The residuals of points to the hypothesis are compared with a threshold to classify points into inliers and outliers at \raisebox{-0.7pt}{5}⃝. The ratio of inliers is utilized as the reward $r_{t}$ for the current action $a_{t}$ to train the policy network $\pi_{\phi}$ at \raisebox{-0.7pt}{6}⃝. Additionally, the action, inliers, and residuals are used for the state transition to generate the next state ${s_{t+1}}$ at \raisebox{-0.7pt}{7}⃝, as shown in Section 3.5. Finally, the environment outputs the next state ${s_{t+1}}$ and the reward $r_{t}$, which are used by the agent to start the next step at \raisebox{-0.7pt}{8}⃝.

Furthermore, the state ${s_{t}}$, action $a_{t}$, reward $r_{t}$, and the next state ${s_{t+1}}$ in every step are collected as experiences in the replay buffer [13] for training. After all episodes are complete, the hypothesis with the highest inlier ratio across all episodes is output as the best hypothesis ${h_{Best}}$ for this pair of images.

The agent receives the state ${s_{t}}\in{\mathbb{R}^{N\times C}}$, where $N$ is the number of the correspondences of a pair of images. The channel $C=c+3$, with $c$ denoting the dimension of data features and $3$ representing the dimension of memory features (see Section 3.5). In addition, the first hypothesis $h_{0}$ of each episode is generated using a randomly sampled minimum set, which provides a basic performance for RLSAC.

Then the state ${s_{t}}$ is fed into the policy network $\pi_{\phi}$. To achieve permutation invariance for the input data, RLSAC uses edge convolution (EdgeConv) from the DGCNN [34] as the basic module to establish the policy network $\pi_{\phi}$. The EdgeConv can model the interrelationship of correspondences by graph neural networks. It can extract features from neighboring nodes in a graph and aggregate them into a central node. The policy network $\pi_{\phi}$ extracts data features and calculates the probabilities by softmax for each point of being in the minimum set. The probability can be interpreted as the probability of obtaining a greater return on the action for the long term, rather than the probability of achieving the best result in the current state. Consequently, the $m$ points with the highest probability are selected as the minimum set $\mathcal{M}_{t}$ in the current state ${s_{t}}$. Importantly, all previously used minimum sets are recorded to avoid selecting a duplicate minimum set. Thus, when a new minimum set is selected, it is checked whether it has been used. If the new minimum set has already been used, then another $m$ data set is selected following the probability, which will be checked as well. Until the unused data set is found, it is output as the minimum set. Ultimately, the minimum set is as the action $a_{t}\in{\mathbb{R}^{m\times C}}$ output by the agent.

The environment generates a hypothesis $h_{t}$ based on the received action $a_{t}$ from the agent. Next, the residuals $R$ are calculated for all data $\chi$ with respect to the hypothesis:

$$R_{t}=\left\{{D\left({{x_{i}},{h_{t}}}\right)\left|{{x_{i}}\in\chi,i=1,2,...,N}\right.}\right\},$$

(5)

where the function $D$ is for residual calculation. The scalar valued residuals $R$ are calculated at all data $\chi$. With the pre-defined threshold, the data $\chi$ is divided into inliers $I_{t}\in{\mathbb{R}^{z\times C}}$ and outliers $O_{t}\in{\mathbb{R}^{\left({N-z}\right)\times C}}$ based to the residuals.

In the sampling consensus methods, the inlier ratio commonly represents the quality of the hypothesis. Since the hypothesis $h_{t}$ is generated by the action $a_{t}$ in RLSAC, the inlier ratio can be considered as the reward $r_{t}$ of the current action $a_{t}$ to train the policy network $\pi_{\phi}$ without ground truth.

With the hypotheses related-rewards, RLSAC can update weights $\phi$ of the policy network $\pi$ without differentiating the sampling process. By modeling sampling consensus as an interactive process in reinforcement learning, RLSAC achieves end-to-end robust estimation while avoiding differentiating.

The state transition module allows RLSAC to encode the previous state and action, enabling RLSAC to explore the state space effectively. As illustrated in Figure 3, the next state $s_{t+1}$ is obtained by concatenating the original data features $N\times c$ with memory features $\left\{{A,R,\mathcal{H}}\right\}$, which comprise action features $N\times 1$, residual features $N\times 1$ and historical features $N\times 1$.

The action features are derived from the current action $a_{t}$. If a data point $x_{i}$ is used in the action $a_{t}$, the corresponding entry in the action features ${\alpha_{i}}$ is set to $1$, otherwise it is $-1$:

$$A_{t}=\left\{{{\alpha_{i}}}\right\},{\alpha_{i}}=\left\{{\begin{array}[]{*{20}{l}}1&{{\rm{if\ }}{x_{i}}\in{a_{t}}}\\ {-1}&{{\rm{otherwise}}}\end{array}}\right.,i=1,2,...,N.$$

(6)

The action features could provide the agent with information on the current action, like the current location in the data for the next state.

The residual features are obtained from the residuals $R_{t}$ computed in Eq. 5. They encode the relative relationship between the hypothesis and the data, which can be thought of as the direction for the agent.

The historical features collect how often is a data point used up to the current step. They are initialized to $0$ and updated as follow: for each data point $x_{i}$ used in the current action $a_{t}$, the corresponding entry in the historical features ${\tau_{i}}$ is incremented by $1$:

$${\mathcal{H}_{t}}=\left\{{{\tau_{i}}}\right\},{\tau_{i}}={\tau_{i}}+1{\rm{\ if\ }}{x_{i}}\in{a_{\rm{t}}},i=1,2,...,N.$$

(7)

The historical features can provide the agent with information on the actions taken so far, enabling it to keep track of its path through the data.

By encoding the state in this way, the agent has access to both current and historical information, allowing it to reach the goal state. Furthermore, the memory features serve as a director of local navigation, guiding the agent to explore the data features and reach its destination gradually.

RLSAC is built on the widely used reinforcement learning framework SAC-Discrete [13]. As shown in Figure 2, the collected replay buffer is retrieved for off-policy to update the actor and critic network in reinforcement learning.

Since RLSAC uses the inlier ratio as the reward for training, it is an unsupervised method. To explore more images, only one episode with multiple steps is collected per image pair during training. The framework is trained for 100 epochs. During training, the episode termination conditions are set as follows: (i) if the number of inliers is unchanged for $\kappa=2$ steps; (ii) if the inlier ratio does not exceed the maximum inlier ratio in the episode for $\varsigma=3$ steps; (iii) if the maximum number of steps $\psi=15$ is reached. During testing, only the third condition is valid, and each pair of images can be tested in $\nu$ episodes.

As with the standard reinforcement learning [13], RLSAC uses probabilistic sampling during training and max sampling during testing. For the EdgeConv module in the policy network, the value of k-nearest neighbors is set to $k=15$. The network architecture and additional information can be found in the supplementary material. Moreover, we have included ablation experiments to analyze the impact of specific network details, also provided in the supplementary material. Similar to the RANSAC implementation, RLSAC also employs refinement using the inliers of the final best hypothesis as the final model polishing.

The experiments are conducted on a Linux computer with an Intel i7 3.6GHz CPU and an NVIDIA RTX 2080Ti GPU. Our implementation is based on PyTorch.

The 2D line fitting task is a basic problem in robust estimation, which can visualize the process of robust estimation and assess the robustness quantitatively to different outlier rates. In this task, each data point contains only the coordinates, thus all data points can be expressed as: $\chi=\left\{{\left[{{x_{i}},{y_{i}}}\right]\left|{\ i=1,2,...,N}\right.}\right\}$, where $\chi$ represents the data features $N\times c$, with $c=2$. Since two points are sufficient to determine a 2D line, the agent can generate a hypothesis by selecting two points from the set of all data points as the minimum set $N\times m$, where $m=2$. For comparison, the RANSAC [16] algorithm is evaluated with the same task.

To evaluate the performance of RLSAC in the 2D line fitting task, we synthesize $N=100$ data points for both training and testing. Specifically, a ground truth 2D line is randomly generated in a $10\times 10$ picture. True inliers are then randomly generated on the line based on the set outlier rate. Next, the true inliers are uniformly disturbed within a range of 0.1 around the line. Thus, the inlier threshold could be set to $\varepsilon=0.1$. Additionally, true outliers are randomly scattered throughout the picture. Moreover, it is possible for true outliers to be located within the inlier region, as is often the case in real-world scenes.

For accurately evaluating and comparing the performance of different methods, we adopt the mean Average Accuracy (mAA) metric in [2]. The angular difference between the estimated line and the ground truth line is used as the error metric, which is used to calculate the mAA metric with a tolerance threshold of $0.5^{\circ}$.

The performance of RLSAC and RANSAC at different outlier rates with 150 iterations is evaluated and the results are presented in Table 1. When the outlier rate is lower than 0.5, RANSAC can achieve a similar performance to RLSAC. This is because the repeated random selection of the minimum set can achieve good performance in simple low outlier rate scenes. However, the performance of RANSAC degrades more quickly than that of RLSAC as the outlier rate increases. This suggests that RLSAC can explore hypotheses closer to the ground truth in noisy scenes.

As the outlier rate increases beyond 0.5, RLSAC can still maintain a low error and high mAA score. This indicates that RLSAC is more robust to disturbances, providing a more stable performance compared to RANSAC.

The qualitative results of RLSAC are shown in Figure 4, which demonstrates its ability to quickly find the better hypothesis even from a bad initial state. This means that the memory features can serve as a director of local navigation, allowing RLSAC to move towards a better state through outputting actions. Moreover, once RLSAC finds a high inlier ratio hypothesis that is close to the ground truth, it continues to explore the local state to further improve the result, rather than randomly transitioning to other states.

The results of the 2D line fitting task show that RLSAC can effectively utilize both data features and memory features, maintaining robust and stable performance in noisy environments. In addition, it can gradually explore a better hypothesis with the help of memory features.

The fundamental matrix estimation is a crucial robust estimation task in computer vision, which solves the fundamental matrix by correspondences in a pair of images. In this study, we use the data and settings in CVPR tutorial RANSAC in 2020 [3], where the correspondences are detected by RootSIFT [1] and matched by the nearest neighbor. The training data comprises 12 scenes, each with 100k image pairs, while the test data includes 2 scenes, with 4950 image pairs. The dataset uses Second Nearest Neighbor (SNN) ratio as the matching score. The correspondences are sorted in descending order by SNN. Then the top $N=150$ SNN correspondences are selected. Next, each point extracts a 128-dimensional descriptor $desc$ through the SIFT algorithm. Thus, the data points are: $\chi=\left\{{\left[{x_{1}^{i},y_{1}^{i},x_{2}^{i},y_{2}^{i},SNN^{i},desc_{1}^{i},desc_{2}^{i}}\right]\left|{i=1,2,...,N}\right.}\right\}\in{\mathbb{R}^{N\times c}}$, where $c=261$. The 8-point solver is used to solve the hypotheses [21]. The inlier threshold of RLSAC is set to $\varepsilon=4$. The evaluation settings follow [3]. To compare with other methods, RANSAC [16], USAC [24], MAGSAC++ [6] are tested, which use the recommended settings in [3].

As shown in Figure 5, RLSAC outperforms other methods in estimating both the rotation matrix and translation vector. Remarkably, RLSAC can achieve high accuracy with only 100 iterations. Comparisons reveals that RLSAC possesses a higher upper bound in performance than MAGSAC.

The quantitative results of the methods at 1k iterations are presented in Table 2. RLSAC achieves great performance with small errors, a little higher than MAGSAC++[6]. Specifically, RLSAC shows high performance in estimating the direction of the translation vector compared to other methods. Since RLSAC could effectively learn features from the data to exclude noise disturbances, it has significant performance on translation vector estimation.

Figure 6 shows the qualitative results of the fundamental matrix estimation of RLSAC in each scene. The figure demonstrates that RLSAC selects rigid and fixed feature points on buildings as the minimum sets, which can help to find better poses. Furthermore, the step-by-step fundamental matrix estimation results of RLSAC are visualized in Figure 7. It demonstrates a gradual increase in the inlier ratio of the hypothesis and a decrease in the rotation and translation errors. This illustrates that RLSAC can progressively explore the state space to find a better hypothesis.