SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jiyao Zhang $^{1,2,3}$ , Mingdong Wu $^{1,3}$ , Hao Dong $^{1,3}$ †
<sup>1</sup> Center on Frontiers of Computing Studies, School of Computer Science, Peking University <sup>2</sup> Beijing Academy of Artificial Intelligence
$^{3}$  National Key Laboratory for Multimedia Information Processing, School of Computer Science, Peking University
jiyaozhang@stu.pku.edu.cn,{wmingd,hao.dong}@pku.edu.cn
Object pose estimation plays a vital role in embodied AI and computer vision, enabling intelligent agents to comprehend and interact with their surroundings. Despite the practicality of category-level pose estimation, current approaches encounter challenges with partially observed point clouds, known as the multihypothesis issue. In this study, we propose a novel solution by reframing category-level object pose estimation as conditional generative modeling, departing from traditional point-to-point regression. Leveraging score-based diffusion models, we estimate object poses by sampling candidates from the diffusion model and aggregating them through a two-step process: filtering out outliers via likelihood estimation and subsequently mean-pooling the remaining candidates. To avoid the costly integration process when estimating the likelihood, we introduce an alternative method that trains an energy-based model from the original score-based model, enabling end-to-end likelihood estimation. Our approach achieves state-of-the-art performance on the REAL275 dataset and demonstrates promising generalizability to novel categories sharing similar symmetric properties without fine-tuning. Furthermore, it can readily adapt to object pose tracking tasks, yielding comparable results to the current state-of-the-art baselines. Our checkpoints and demonstrations can be found at https://sites.google.com/view/genpose.
Object pose estimation is a fundamental problem in the domains of robotics and machine vision, providing a high-level representation of the surrounding world. It plays a crucial role in enabling intelligent agents to understand and interact with their environments, supporting tasks such as robotic manipulation, augmented reality, and virtual reality [1; 2; 3]. In the existing literature, category-level object pose estimation [4; 5; 6] has emerged as a more practical approach compared to instance-level pose estimation since the former eliminates the need for a 3D CAD model for each individual instance. The main challenge of this task is to capture the general properties while accommodating the intra-class variations, i.e., variations among different instances within a category [7; 5; 8].
Figure 1: Multi-hypothesis issue in category-level object pose estimation comes from (a) symmetric objects and (b) partial observation.
*: equal contribution, †: corresponding author
To address this task, previous studies have explored regression-based approaches. These approaches involve either recovering the object pose by predicting correspondences between the observed point cloud and synthesized canonical coordinates [4; 7], or estimating the pose in an end-to-end manner [9; 6]. The former approaches leverage category prior information, such as the mean point cloud, to handle intra-class variations. As a result, they have demonstrated strong performance across various metrics. On the other hand, the latter approaches do not rely on the coordinate mapping process, which inherently makes them faster and amenable to end-to-end training.
Despite the promising results of previous approaches, a fundamental issue inherent to regression-based training persists. This issue, known as the multi-hypothesis problem, arises when dealing with a partially observed point cloud. In such cases:
Multiple feasible pose hypotheses can exist, but the network can only be supervised for a single pose due to the regression-based training.
Figure 1(a) illustrates this problem by showcasing the feasible ground truth poses of symmetric objects (e.g., bowls). Moreover, partially observed point clouds of an object may appear similar from certain views (e.g., a mug with an obstructed handle may look the same from certain views, as shown in Figure 1(b)), further exacerbating the multiple hypothesis issue. Previous studies have proposed ad-hoc solutions to tackle this issue, such as designing special network architectures [10; 11] or augmenting the ground truth poses [4; 5] for symmetric objects. However, these approaches cannot fundamentally resolve this issue due to the lack of generality. Consequently, an ideal object pose estimator should model a pose distribution instead of regressing a single pose when presented with a partially observed point cloud.
To this end, we formulate category-level object pose estimation as conditional generative modeling that inherently models the distribution of multiple pose hypotheses conditioned on a partially observed point cloud. To model the conditional pose distribution, we employ score-based diffusion models, which have demonstrated promising results in various conditional generation tasks [12]. Specifically, we estimate the score functions, i.e., the gradients of the log density, of the conditional pose distribution perturbed by different noise levels. These trained gradient fields can provide guidance to iteratively refine a pose, making it more compatible with the observed point cloud. In the testing phase, we estimate the object pose of a partial point cloud by sampling from the conditional pose distribution using an MCMC process incorporated with the learned gradient fields. As a result, our method can propose multiple pose hypotheses for a partial point cloud, thanks to stochastic sampling.
However, it is essential to note that a sampled pose might be an outlier of the conditional pose distribution, meaning it has a low likelihood, which can severely hurt the performance. To address this issue, we sample a group of pose candidates and then filter out the outliers based on the estimated likelihoods. Unfortunately, estimating the likelihoods from score-based models requires a highly time-consuming integration process [13], rendering this approach impractical. To overcome this challenge, we propose an alternative solution that trains an energy-based model from the original score-based model for likelihood estimation. Following training, the energy network is guaranteed to estimate the log-likelihood of the original data distribution up to a constant. By ranking the candidates according to the energy network's output, we can filter out candidates with low ranks (e.g., the last  $40\%$ ). Finally, the output pose is aggregated by mean-pooling the remaining pose candidates.
Our experiments demonstrate several superiorities over the existing methods. Firstly, without any ad-hoc network or loss design for symmetry objects, our method achieves state-of-the-art performance and surpasses  $50\%$  and  $60\%$  on the strict  $5^{\circ}2\mathrm{cm}$  and  $5^{\circ}5\mathrm{cm}$  metrics, respectively, on the REAL275 dataset. Besides, our method can directly generalize to novel categories without any fine-tuning to some degree due to its prior-free nature. Our methods can be easily adapted to the object pose tracking task with few modifications and achieve comparable performance with the SOTA baselines.
- We study a fundamental problem, namely multi-hypothesis issue, that has existed in category-level object pose estimation for a long time and introduce a generative approach to address this issue.
- We propose a novel framework that leverages the energy-based diffusion model to aggregate the candidates generated by the score-based diffusion model for object pose estimation.
- Our framework achieves exceptionally state-of-the-art performance on existing benchmarks, especially on symmetric objects. In addition, our framework is capable of object pose tracking and can generalize to objects from novel categories sharing similar symmetric properties.
2.1 Category Level Object Pose Estimation
Category-level object pose estimation [4; 14; 5; 15] aims to predict pose of various instances within the same object category without requiring CAD models. To this end, NOCS [4] introduce normalized object coordinate space for each object category and predicts the corresponding object shape in canonical space. Object pose is then calculated by pose fitting using the Umayama [16] algorithm. CASS [17], DualPoseNet [9], and SSP-Pose [18] end-to-end regress the pose and simultaneously reconstruct the object shape in canonical space to improve the precision of pose regression. SAR-Net [10] and GPV-Pose [11] specifically design symmetric-aware reconstruction to enhance the performance of predicting the translation and size of symmetrical objects while regressing the object pose. However, direct regression of pose or object coordinates in canonical space fail when there is considerable intra-class variation. To overcome this, FS-Net [6] introduces an online box-cage-based 3D deformation augmentation method, and RBP-Pose [19] proposes a nonlinear data augmentation method. On the other hand, SPD [7], SGPA [5], CR-Net [20], and DPDN [8] introduce category-level object priors (e.g., mean shape) and learn how to deform the prior to get the object coordinates in canonical space. These methods have demonstrated strong performance. Overall, the existing category-level object pose estimation methods are regression-based. However, considering symmetric objects (e.g., bowls, bottles, and cans) and partially observed objects (e.g., mugs without visible handles), there are multiple plausible hypotheses for object pose. Therefore, the problem of object pose estimation should be formulated as a generation problem rather than a regression problem.
2.2 Score-based Generative Models
In the realm of estimating the gradient of the log-likelihood pertaining to specified data distribution, a pioneering approach known as the score-based generative model [21; 22; 23; 24; 25; 13; 26] was originally introduced by [22]. In an effort to provide a feasible substitute objective for score-matching, the denoising score-matching (DSM) technique[21] has further put forward a viable proposition. To enhance the scalability of the score-based generative model, [23] has introduced a sliced score-matching objective, which involves projecting the scores onto random vectors prior to their comparison. They have also presented annealed training for denoising score matching[24] and have introduced improved training techniques to complement these methods [25]. Additionally, they have extended the discrete levels of annealed score matching to a continuous diffusion process, thereby demonstrating promising results in the domain of image generation [13]. Recent studies have delved further into exploring the design choices of the diffusion process [27], maximum likelihood training [26], and deployment on the Riemann manifold [28]. These advancements have showcased encouraging outcomes when applying score-based generative models to high-dimensional domains, thus fostering their broad utilization across various fields, such as object rearrangement [29], medical imaging [30], point cloud generation [31], scene graph generation [32], point cloud denoising [33], depth completion [34], and human pose estimation [35]. These studies have formulated perception-related problems as either conditional generative modeling or in-painting tasks, thereby harnessing the power of score-based generative models to tackle these challenges. In contrast to these approaches, we incorporate the score-based diffusion model with an additional energy-based diffusion model [36; 37] for filtering out the outliers so as to improve the performance by aggregating the remaining candidates.
Task Description: In this work, we aim to estimate 6D object pose from the partially observed point cloud. The learning agent is given a training set with paired object poses and point clouds  $\mathcal{D} = \{(\pmb {p}_i,O_i)\}_{i = 1}^n$ , where  $\pmb {p}_i\in \mathrm{SE}(3)$  and  $O_{i}\in \mathbb{R}^{3\times N}$  denote a 6D pose and a partially observed 3D point cloud with  $N$  points respectively. Given an unseen point cloud  $O^{*}$ , the goal is to recover the corresponding ground-truth pose  $\pmb{p}^{*}$ .
Overview: We formulate the object pose estimation as a conditional generative modeling problem and train a score-based diffusion model  $\Phi_{\theta}$  using the dataset  $\mathcal{D}$ . Further, an energy-based diffusion model  $\Psi_{\phi}$  is trained from the score-based model  $\Phi_{\theta}$  for likelihood estimation. During test time, given an unseen point cloud  $O^{*}$ , we first sample a group of pose candidates  $\{\hat{p}_1,\hat{p}_2,\dots \hat{p}_K\}$  via the score-based diffusion model  $\Phi_{\theta}$ . Then we sort the candidates  $\{\hat{p}_i\}_{i = 1}^K$  in descending order according
(I) Training
(II) Inference
Figure 2: Overview. (I) A score-based diffusion model  $\Phi_{\theta}$  and an energy-based diffusion model  $\Psi_{\phi}$  is trained via denoising score-matching. (II) a) We first generate pose candidates  $\{\hat{p}_i\}_{i=1}^K$  from the score-based model and then b) compute the pose energies  $\Psi_{\phi}(\hat{p}_i, \epsilon|O^*)$  for candidates via the energy-based model. c) Finally, we rank the candidates with the energies and then filter out low-ranking candidates. The remaining candidates are aggregated into the final output by mean-pooling.
to the energy outputs and filter out the last  $1 - \delta \%$  candidates (e.g.,  $\delta = 60\%$ ). Finally, we aggregate the remaining candidates by mean-pooling to obtain the estimated pose  $\hat{\pmb{p}}$ .
3.1 Sampling Pose Candidates via Score-based Diffusion Model
To address the multi-hypothesis issue, an ideal pose estimator should be capable of being trained on multiple feasible ground truth poses given the same point cloud. Also, the ideal pose estimator should be able to output all possible pose hypotheses of the given point cloud during inference.
To this end, we propose to tackle object pose estimation in a conditional generative modeling paradigm. We assume dataset  $\mathcal{D}$  is sampled from an implicit joint distribution  $\mathcal{D} = \{(p_i,O_i)\sim p_{\mathrm{data}}(p,O)\}$ . We aim to model the conditional pose distribution  $p_{\mathrm{data}}(p|O)$  during training, and sample pose hypotheses of an unseen point cloud  $O^{*}$  from  $p_{\mathrm{data}}(p|O^{*})$  during test-time.
Specifically, we employ a score-based diffusion model [24; 23; 13] to estimate the conditional distribution  $p_{\mathrm{data}}(\pmb{p} | O)$ . We adopt Variance-Exploding (VE) Stochastic Differential Equation (SDE) proposed by [13] to construct a continuous diffusion process  $\{\pmb{p}(t)\}_{t=0}^{1}$  indexed by a time variable  $t \in [0,1]$  where  $\pmb{p}(0) \sim p_{\mathrm{data}}(\pmb{p} | O)$  denotes the ground truth pose of the point cloud  $O$ . As the  $t$  increases from 0 to 1, the time-indexed pose variable  $\pmb{p}(t)$  is perturbed by the following SDE:
$$
d \boldsymbol {p} = \sqrt {\frac {d [ \sigma^ {2} (t) ]}{d t}} d \mathbf {w}, \sigma (t) = \sigma_ {\min } \left(\frac {\sigma_ {\max }}{\sigma_ {\min }}\right) ^ {t} \tag {1}
$$
where  $\sigma_{\mathrm{min}} = 0.01$  and  $\sigma_{\mathrm{max}} = 50$  are hyper-parameters.
During training, we aim to estimate the score function of the perturbed conditional pose distribution  $\nabla_{\pmb{p}}\log p_t(\pmb {p}|O)$  of all  $t$  , where the  $p_t(p|O)$  denotes the marginal distribution of  $\pmb {p}(t)$  ..
$$
p _ {t} (\boldsymbol {p} (t) | O) = \int \mathcal {N} (\boldsymbol {p} (t); \boldsymbol {p} (0), \sigma^ {2} (t) \mathbf {I}) \cdot p _ {0} (\boldsymbol {p} (0) | O) d \boldsymbol {p} (0) \tag {2}
$$
Notably, when  $t = 0$ ,  $p_0(\pmb{p}(0)|O) = p_{\mathrm{data}}(\pmb{p}(0)|O)$  is exactly the data distribution.
Thanks to the Denoising Score Matching (DSM) [21], we can obtain a guaranteed estimation of  $\nabla_{\pmb{p}} p_t(\pmb{p} | O)$  by training a score network  $\Phi_\theta: \mathbb{R}^{|\mathcal{P}|} \times \mathbb{R}^1 \times \mathbb{R}^{3 \times N} \to \mathbb{R}^{|\mathcal{P}|}$  via the following objective:
$$
\mathcal {L} (\theta) = \mathbb {E} _ {t \sim \mathcal {U} (\epsilon , 1)} \left\{\lambda (t) \mathbb {E} _ {\substack {\boldsymbol {p} (0) \sim p _ {\text {data}} (\boldsymbol {p} (0) | O), \\ \boldsymbol {p} (t) \sim \mathcal {N} (\boldsymbol {p} (t); \boldsymbol {p} (0), \sigma^ {2} (t) \mathbf {I})}} \left[ \left\| \boldsymbol {\Phi} _ {\theta} (\boldsymbol {p} (t), t | O) - \frac {\boldsymbol {p} (0) - \boldsymbol {p} (t)}{\sigma (t) ^ {2}} \right\| _ {2} ^ {2} \right] \right\} \tag{3}
$$
where  $\epsilon$  is a hyper-parameter that denotes the minimal noise level. When minimizes the objective in Eq. 3, the optimal score network satisfies  $\Phi_{\theta}(\pmb{p}, t|O) = \nabla_{\pmb{p}} \log p_t(\pmb{p}|O)$  according to [21].
After training, we can approximately sample pose candidates  $\{\hat{\pmb{p}}_i\}_{i = 1}^K$  from  $p_{\mathrm{data}}(\pmb {p}|O)$  by sampling from  $p_{\epsilon}(\pmb {p}|O)$ , as  $\lim_{\epsilon \to 0}p_{\epsilon}(\pmb {p}|O) = p_{\mathrm{data}}(\pmb {p}|O)$ . To sample from  $p_{\epsilon}(\overline{\pmb{p}} |O)$ , we can solve the following Probability Flow (PF) ODE [13] where  $\pmb {p}(1)\sim \mathcal{N}(\mathbf{0},\sigma_{\max}^2\mathbf{I})$ , from  $t = 1$  to  $t = \epsilon$ :
$$
\frac {d \boldsymbol {p}}{d t} = - \sigma (t) \dot {\sigma} (t) \nabla_ {\boldsymbol {p}} \log p _ {t} (\boldsymbol {p} | O) \tag {4}
$$
where the score function  $\log p_t(\pmb{p} | O)$  is empirically approximated by the estimated score network  $\Phi_{\theta}(\pmb{p}, t | O)$  and the ODE trajectory is solved by RK45 ODE solver [38].
In practice, the score network  $\Phi_{\theta}(\pmb{p}, t|O)$  is implemented without any ad-hoc design for the symmetric objects: The point cloud  $O$  is encoded into a global feature by PointNet++ [39], the input pose  $\pmb{p}$  is encoded by feed-forward MLPs and the time variable  $t$  is encoded by a commonly used projection layer following [13]. The pose  $\pmb{p}$  is represented as a 9-D variable  $[R|T]$ , where  $R \in \mathbb{R}^6$  and  $T \in \mathbb{R}^3$  denote rotation and translation vectors, respectively. Due to the discontinuity of quaternions and Euler angles in Euclidean space, we employ the continuous 6-D rotation representation  $[R_x|R_y]$  following[6; 40]. We defer full details into Appendix A.
3.2 Aggregating Pose Candidates via Energy-based Diffusion Model
Though we can sample pose candidates  $\{\hat{p}_i\}_{i=1}^K$  from the conditional pose distribution  $p_{\mathrm{data}}(p|O)$  via the learned score model  $\Phi_\theta$ , we still need to determine a final output estimation  $\hat{p} \in \mathrm{SE}(3)$ . In other words, the pose candidates need to be aggregated into a single estimation.
An initial approach to aggregating candidates is mean pooling. However, when sampling candidates from  $p_{\mathrm{data}}(\pmb{p}|O)$ , there is no guarantee that the sampled candidates will have a high likelihood. This means that outlier poses in low-density regions are still likely to be included in the mean-pooled pose, which can negatively impact its performance. Therefore, we propose estimating the data likelihoods  $\{p_{\mathrm{data}}(\hat{\pmb{p}}_i|O)\}_{i=1}^K$  to rank the candidates and then filter out low-ranking candidates.
Unfortunately, estimating likelihood via the score model itself requires a time-consuming integration process [13], which is impractical for real-time applications:
$$
\log p _ {\mathrm {d a t a}} (\boldsymbol {p} | O) \approx \log p _ {\epsilon} (\boldsymbol {p} | O) = \log p _ {1} (\boldsymbol {p} (1) | O) - \frac {1}{2} \int_ {\epsilon} ^ {1} \frac {d [ \sigma^ {2} (t) ]}{d t} \boldsymbol {\nabla} \cdot \boldsymbol {\Phi} _ {\theta} (\boldsymbol {p} (t), t | O) d t \tag {5}
$$
where  $\{\pmb{p}(t)\}_{t\in [0,1]}$  is the same forward diffusion process as in Eq. 1.
To this end, we propose to train an energy-based model  $\Psi_{\phi}:\mathbb{R}^{|\mathcal{P}|}\times \mathbb{R}^1\times \mathbb{R}^{3\times N}\to \mathbb{R}^1$  that enables an end-to-end surrogate estimation of the data likelihood by supervising the energy-induced gradient  $\nabla p\Psi_{\phi}(\pmb {p},t|O)$  with the denoising score-matching objective in Eq. 3:
$$
\mathcal {L} (\phi) = \mathbb {E} _ {t \sim \mathcal {U} (\epsilon , 1)} \left\{\lambda (t) \mathbb {E} _ {\substack {\boldsymbol {p} (0) \sim p _ {0} (\boldsymbol {p} | O), \\ \boldsymbol {p} (t) \sim p _ {0 t} (\boldsymbol {p} (t) | \boldsymbol {p} (0), O)}} \left[ \left\| \nabla_ {\boldsymbol {p} (t)} \boldsymbol {\Psi} _ {\phi} (\boldsymbol {p} (t), t | O) - \frac {\boldsymbol {p} (0) - \boldsymbol {p} (t)}{\sigma (t) ^ {2}} \right\| _ {2} ^ {2} \right] \right\} \tag{6}
$$
In this way, the optimal energy model holds  $\nabla_{\pmb{p}}\Psi_{\phi}^{*}(\pmb{p},t|O) = \nabla_{\pmb{p}}\log p_t(\pmb{p}|O)$ , or equivalently,  $\Psi_{\phi}^{*}(\pmb{p},t|O) = \log p_{t}(\pmb{p}|O) + C$  where  $C$  is a constant. Although the optimal energy model and the ground truth likelihood differ by a constant  $C$ , this does not prevent it from functioning as a good surrogate likelihood estimator for ranking the candidates:
$$
\Psi_ {\phi} ^ {*} (\boldsymbol {p} _ {i}, \epsilon | O) > \Psi_ {\phi} ^ {*} (\boldsymbol {p} _ {j}, \epsilon | O) \iff \log p _ {\epsilon} (\boldsymbol {p} _ {i} | O) > \log p _ {\epsilon} (\boldsymbol {p} _ {j} | O) \tag {7}
$$
Nevertheless, training an energy-based diffusion model from Eq. 6 is known to be difficult and time-inefficient due to the need to calculate the second-order derivations in the objective function. Following [36], we parameterize the energy-based model as  $\Psi_{\phi}(\pmb{p}, t|O) = \langle \pmb{p}, \Phi_{\phi}(\pmb{p}, t|O) \rangle$  to alleviate the training burden. We defer the full training details to Appendix A.
With the trained energy model, we sort the candidates into a sequence  $\hat{\pmb{p}}_{\tau_1} \succ \hat{\pmb{p}}_{\tau_2} \dots \succ \hat{\pmb{p}}_{\tau_K}$  where:
$$
\hat {\boldsymbol {p}} _ {\tau_ {i}} \succ \hat {\boldsymbol {p}} _ {\tau_ {j}} \Longleftrightarrow \Psi_ {\phi} (\hat {\boldsymbol {p}} _ {\tau_ {i}}, \epsilon | O) > \Psi_ {\phi} (\hat {\boldsymbol {p}} _ {\tau_ {j}}, \epsilon | O) \tag {8}
$$
Subsequently, we filter out the last  $1 - \delta \%$  candidates and obtain  $\hat{\pmb{p}}_{\tau_1} \succ \hat{\pmb{p}}_{\tau_2} \dots \succ \hat{\pmb{p}}_{\tau_M}$  where  $\delta \in (0,1)$  is a hyper parameter and  $M = \lfloor \delta \cdot K \rfloor$ .
Finally, we aggregate the remaining candidates  $\{\hat{p}_{\tau_i} = (\hat{T}_{\tau_i},\hat{R}_{\tau_i})\}_{i = 1}^M$  by averaging the rotations  $\{\hat{R}_{\tau_i}\}_{i = 1}^M$  and the translations  $\{\hat{T}_{\tau_i}\}_{i = 1}^M$  respectively, to obtain the output pose  $\hat{\pmb{p}} = (\hat{T},\hat{R})$ . In specific, the translations are pooled by vanilla averaging  $\hat{T} = \frac{\sum_{i = 1}^{M}\hat{T}_{\tau_i}}{M}$ . To obtain the averaged rotation, we initially translate the rotations into quaternions  $\{\hat{q}_{\tau_i}\}_{i = 1}^M$ . Following [41], the average quaternion can then be found by the following maximization procedure:
$$
\hat {\boldsymbol {q}} = \underset {\boldsymbol {q} \in S O (3)} {\arg \max } \boldsymbol {q} ^ {T} \left(\frac {\sum_ {i = 1} ^ {M} A \left(\hat {\boldsymbol {q}} _ {\tau_ {i}}\right)}{M}\right) \boldsymbol {q}, A \left(\hat {\boldsymbol {q}} _ {\tau_ {i}}\right) = \hat {\boldsymbol {q}} _ {\tau_ {i}} \hat {\boldsymbol {q}} _ {\tau_ {i}} ^ {T} \tag {9}
$$
By definition, the solution of Eq. 9 is the eigenvector of the  $4 \times 4$  matrix  $\frac{\sum_{i=1}^{M} A(\mathbf{q}_{\tau_i})}{M}$  corresponding to the maximum eigenvalue, which can be efficiently solved by QUEST algorithm [42].
Despite addressing the multi-hypothesis issue, our method offers several additional advantages: