---
title: Cross-Entropy Method with Reward Model
url: https://www.emergentmind.com/topics/cross-entropy-method-with-reward-model-cem-rm
type: topic
---

# Cross-Entropy Method with Reward Model

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 [2109.06619, 1909.12830].

## 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{T}(s, \mathbf{a})$ yields the transformed source, $\tilde{x}_i = R(\mathbf{e}) x_i + \mathbf{t}$. 

The reward $R(s, \mathbf{a})$ is designed to penalize misalignment (e.g., partial overlap), and is often set as the negative of a maximum-consensus style alignment error $D^{mc}(\tilde{X}, Y)$, such that $R(s, \mathbf{a}) = -D^{mc}(\tilde{X}, Y)$ [2109.06619].

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

The alignment error $D^{mc}$ incorporates inlier recognition under a spatial threshold $\varepsilon$. For each transformed source point $\tilde{x}_i$, the minimum distance to target points is evaluated, $d_{\tilde{x}_i, Y} = \min_j \|\tilde{x}_i - y_j\|_2$. Inliers are those with $d_{\tilde{x}_i, Y} \leq \varepsilon$. 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 $\rho_\varepsilon(u) = \mathbf{1}\{u \leq \varepsilon\}$.

To incorporate look-ahead, CEM-RM augments each candidate’s current reward $R(s, \mathbf{a})$ with a “future” reward $R(s', \mathbf{a}_\mathrm{ICP})$, where $s' = \mathcal{T}(s, \mathbf{a})$ and $\mathbf{a}_\mathrm{ICP}$ is a corrective update derived from running one round of ICP. The fused score is
$$
S(s, \mathbf{a}) = \alpha R(s, \mathbf{a}) + (1-\alpha) R(s', \mathbf{a}_\mathrm{ICP}), \quad 0\leq \alpha \leq 1 .
$$
This approach encourages selection of transformations that are both immediately rewarding and admit efficient fine registration [2109.06619].

## 3. Differentiable Cross-Entropy Method Module

CEM-RM employs a differentiable variant of CEM (“DCEM”; see [1909.12830]) to enable end-to-end training with gradient-based methods. The classical CEM updates the sampling distribution’s parameters (mean $\mu_t$, variance $\sigma^2_t$) by maximum-likelihood estimation on a hard top-$k$ 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-$k$ projection ($\mathrm{sparsemax}(q)$), which maps the fused scores $q_1, \dots, q_N$ to a sparse probability simplex vector $\mathbf{p}$. The update equations are
$$
\mu_{t+1} = \sum_i p_i \mathbf{a}_t^i, \quad \sigma^2_{t+1} = \sum_i p_i (\mathbf{a}_t^i - \mu_{t+1})^2 \,.
$$
The algorithm employs the reparameterization trick for low-variance gradients. All operations are differentiable, supporting gradient flow from the output mean $\hat{a} = \mu_T$ back through network weights [1909.12830].

**Pseudocode Sketch: Differentiable CEM-RM Core Loop**

```python
# Inputs: s = {X, Y}, prior_net P_w
mu, sigma = P_w(s)
for t in range(T):
    for i in range(N):
        eps = sample_normal()
        a_i = mu + sigma * eps
        s_i = T(s, a_i)
        r_i = R(s, a_i)
        if t < M:
            a_icp, r_fut = ICP_and_future_reward(s_i)
            q_i = alpha * r_i + (1 - alpha) * r_fut
        else:
            q_i = r_i
    p = sparsemax(q)
    mu = sum(p_i * a_i for i in range(N))
    sigma = sqrt(sum(p_i * (a_i - mu)**2 for i in range(N)))
```

## 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, $\mathcal{N}(\mu_0, \sigma_0^2)$. The prior network is a PointNet-style architecture. Features $F_X$ and $F_Y$ are computed per point (via DGCNN), soft matching is derived for correspondences, and SVD recovers rotation/translation to populate $\mu_0$. The standard deviation $\sigma_0$ 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 [2109.06619].

## 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, 
$$
\rho_\mu(u) = \frac{\mu u^2}{\mu + u^2}
$$ 
as a metric for registration quality over all inlier distances,
$$
\mathcal{L}(w) = \mathbb{E}_{(X, Y)}\left[ \frac{1}{N}\sum_{i=1}^N \rho_\mu(d_{\tilde{x}_i, Y}) + \frac{1}{M}\sum_{j=1}^M \rho_\mu(d_{y_j, \tilde{X}}) \right]
$$
where $\tilde{X}$ is the source cloud transformed by the predicted $\hat{a} = \mu_T$. This penalizes small distances quadratically while treating outliers sublinearly, promoting robust convergence [2109.06619].

## 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 [2109.06619, 1909.12830]. 

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 [1909.12830]. 

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-$k$ 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 [2109.06619, 1909.12830].

Source: https://www.emergentmind.com/topics/cross-entropy-method-with-reward-model-cem-rm