Cross-Entropy Method with Reward Model

• CEM-RM is an end-to-end differentiable framework that integrates cross-entropy sampling with a learned reward model for optimizing problems like 3D point cloud registration.
• It employs a differentiable cross-entropy method using soft elite selection (sparsemax) to enable gradient flow and robust parameter updates.
• The framework fuses current and future rewards with a data-driven prior to guide the search in non-convex optimization tasks such as camera pose refinement.

The Cross-Entropy Method with Reward Model (CEM-RM) integrates model-based stochastic search with model-guided reward evaluation in a differentiable, end-to-end pipeline. CEM-RM extends the classical cross-entropy method (CEM) by coupling it with a learned or structured reward model, using soft differentiable elite selection and enabling gradient flow through the optimization procedure. This framework supports complex, non-convex optimization tasks such as unsupervised 3D point cloud registration, where candidate solutions are iteratively refined with both current and look-ahead (“future”) reward feedback, and where the search is initialized with a data-driven, learned prior over the transformation space (Jiang et al., 2021, Amos et al., 2019).

1. Markov Decision Process Formulation

CEM-RM models the target application, such as 3D point cloud registration, as a Markov Decision Process (MDP) $(\mathcal{S},\mathcal{A},\mathcal{T},R)$. The state space $\mathcal{S}$ contains tuple pairs of point clouds $s = \{X, Y\}$, with $X = \{x_i \in \mathbb{R}^3\}_{i=1}^N$ and $Y = \{y_j \in \mathbb{R}^3\}_{j=1}^M$. The action space $\mathcal{A}$ is defined as rigid transformations in $SE(3)$ parameterized by Euler angles $\mathbf{e} \in [-\pi, \pi]^3$ and translation $\mathbf{t} \in \mathbb{R}^3$, so that $$\mathbf{a} = [\mathbf{e}; \mathbf{t}] \in \mathbb{R}^6$$. The state transition $\mathcal{S}$0 yields the transformed source, $\mathcal{S}$1.

The reward $\mathcal{S}$2 is designed to penalize misalignment (e.g., partial overlap), and is often set as the negative of a maximum-consensus style alignment error $\mathcal{S}$3, such that $\mathcal{S}$4 (Jiang et al., 2021).

2. Reward Modeling: Maximum-Consensus Criterion and Fused Future Score

The alignment error $\mathcal{S}$5 incorporates inlier recognition under a spatial threshold $\mathcal{S}$6. For each transformed source point $\mathcal{S}$7, the minimum distance to target points is evaluated, $\mathcal{S}$8. Inliers are those with $\mathcal{S}$9. The maximum-consensus metric is: \begin{align*} D^{mc}(\tilde{X}, Y) = 2 &- \frac{1}{N} \sum_{i=1}^N \rho_\varepsilon(d_{\tilde{x}i, Y}) \left(1 - \frac{d{\tilde{x}i, Y}}{\varepsilon}\right) \ &- \frac{1}{M} \sum{j=1}^M \rho_\varepsilon(d_{y_j, \tilde{X}}) \left(1 - \frac{d_{y_j, \tilde{X}}}{\varepsilon}\right) \end{align*} where $s = \{X, Y\}$0.

To incorporate look-ahead, CEM-RM augments each candidate’s current reward $s = \{X, Y\}$1 with a “future” reward $s = \{X, Y\}$2, where $s = \{X, Y\}$3 and $s = \{X, Y\}$4 is a corrective update derived from running one round of ICP. The fused score is

$s = \{X, Y\}$5

This approach encourages selection of transformations that are both immediately rewarding and admit efficient fine registration (Jiang et al., 2021).

3. Differentiable Cross-Entropy Method Module

CEM-RM employs a differentiable variant of CEM (“DCEM”; see (Amos et al., 2019)) to enable end-to-end training with gradient-based methods. The classical CEM updates the sampling distribution’s parameters (mean $s = \{X, Y\}$6, variance $s = \{X, Y\}$7) by maximum-likelihood estimation on a hard top-$s = \{X, Y\}$8 set of elite samples. This is non-differentiable.

DCEM achieves differentiability by replacing the hard selection with a soft selection using sparsemax or soft top-$s = \{X, Y\}$9 projection ($X = \{x_i \in \mathbb{R}^3\}_{i=1}^N$0), which maps the fused scores $X = \{x_i \in \mathbb{R}^3\}_{i=1}^N$1 to a sparse probability simplex vector $X = \{x_i \in \mathbb{R}^3\}_{i=1}^N$2. The update equations are

$X = \{x_i \in \mathbb{R}^3\}_{i=1}^N$3

The algorithm employs the reparameterization trick for low-variance gradients. All operations are differentiable, supporting gradient flow from the output mean $X = \{x_i \in \mathbb{R}^3\}_{i=1}^N$4 back through network weights (Amos et al., 2019).

Pseudocode Sketch: Differentiable CEM-RM Core Loop

$Y = \{y_j \in \mathbb{R}^3\}_{j=1}^M$5

4. Sampling Network Prior: Data-Driven Initialization

To bias the search distribution towards promising regions in transformation space, CEM-RM introduces a learned Gaussian prior over actions, $X = \{x_i \in \mathbb{R}^3\}_{i=1}^N$5. The prior network is a PointNet-style architecture. Features $X = \{x_i \in \mathbb{R}^3\}_{i=1}^N$6 and $X = \{x_i \in \mathbb{R}^3\}_{i=1}^N$7 are computed per point (via DGCNN), soft matching is derived for correspondences, and SVD recovers rotation/translation to populate $X = \{x_i \in \mathbb{R}^3\}_{i=1}^N$8. The standard deviation $X = \{x_i \in \mathbb{R}^3\}_{i=1}^N$9 is learned by max-pooling features to a global vector, followed by an MLP with sigmoid activation. This network is co-trained with the main CEM-RM pipeline (Jiang et al., 2021).

5. Robust Loss: Geman–McClure Estimator

Ground-truth transformations are typically unavailable in unsupervised registration, motivating a robust, unsupervised alignment loss. CEM-RM adopts a scaled Geman–McClure penalty, $Y = \{y_j \in \mathbb{R}^3\}_{j=1}^M$0 as a metric for registration quality over all inlier distances,

$Y = \{y_j \in \mathbb{R}^3\}_{j=1}^M$1

where $Y = \{y_j \in \mathbb{R}^3\}_{j=1}^M$2 is the source cloud transformed by the predicted $Y = \{y_j \in \mathbb{R}^3\}_{j=1}^M$3. This penalizes small distances quadratically while treating outliers sublinearly, promoting robust convergence (Jiang et al., 2021).

6. End-to-End Training and Generalization

Backpropagation proceeds through the full CEM-RM stack: gradients flow from the robust loss through the differentiable CEM module (including sparsemax and reparameterization), influencing the prior network and, consequently, the sampling distribution. During inference, the prior network produces a suitable initialization, and CEM iteratively refines the estimate guided by current and look-ahead rewards (Jiang et al., 2021, Amos et al., 2019).

The modularity of this approach enables generalization to any continuous state-action problem with computable reward and accessible local optimizer for future estimates. Examples include camera pose refinement (using, e.g., photometric or geometric alignment rewards) and simulation-based control tasks with look-ahead estimation by trajectory unrolling.

7. Connections and Extensions

CEM-RM extends prior work on the cross-entropy method and differentiable CEM. The differentiable mechanisms allow integrating learned neural reward models and enable meta-learning of reward structures and priors (Amos et al., 2019).

A key distinction is the integration of a fused reward (combining current and predicted future reward) for sample evaluation, which is not present in classical CEM or standard DCEM formulations. The introduction of sparsemax-based elite selection maintains the search’s exploratory diversity while preserving differentiability, mitigating the gradient estimation issues found in hard top-$Y = \{y_j \in \mathbb{R}^3\}_{j=1}^M$4 selection.

Potential extensions include hybridization with other differentiable sampling-based optimizers (e.g., CMA-ES), meta-learning of hyperparameters (sample count, elite fraction, temperature), and compositional reward modeling for more complex tasks. The CEM-RM framework thus provides a flexible, end-to-end differentiable scheme uniting sampling-based optimization and reward modeling in domains where gradient-based search and robust evaluation are essential (Jiang et al., 2021, Amos et al., 2019).