Reinforcement learning with Mixture-of-Experts (MoE) policies is inherently a latent-variable problem: at each timestep the policy samples a latent route $z_t$ and then an action $a_t$, and that routing sequence affects all future states. If we discard the sampled routing path and try to train only on the marginal mixture over experts, the natural approximations give rise to a nested Monte Carlo estimator in the sense of Rainforth et al. (2018)¹, which has provably poor statistical properties. The central claim of this note is that to avoid this nested-estimator pathology, one must perform routing replay: log the sampled routes and train on the resulting complete-data objective.

The remainder of the post unpacks this statement: we first formalize MoE RL as a latent-variable model (Sections 1–2), then characterize the nested estimator that arises when routes are marginalized and approximated (Section 4.1), and finally explain why routing replay restores a single-level, well-behaved estimator.

---

1. MoE policy as a latent-variable model

We consider an MoE policy in an RL setting with state $s_t$, action $a_t$, and reward $r_t$ at time $t$. In addition to these observed variables, the policy includes a latent routing variable $z_t \in \{1,\dots,K\}$ indicating which expert is selected at that timestep.

The router (gating network) defines a distribution over experts $\rho_\phi(z_t \mid s_t)$, and the expert policies define conditional action distributions $\pi_\theta(a_t \mid s_t, z_t)$.

The joint policy over $(a_t, z_t)$ is

$$\pi_{\theta,\phi}(a_t, z_t \mid s_t) = \rho_\phi(z_t \mid s_t)\,\pi_{\theta}(a_t \mid s_t, z_t).$$

The marginal action policy (what the environment actually sees) is the MoE mixture:

$$\bar{\pi}_{\theta,\phi}(a_t \mid s_t) = \sum_{z_t} \rho_\phi(z_t \mid s_t)\,\pi_{\theta}(a_t \mid s_t, z_t).$$

So $z_t$ is literally a latent variable in a mixture model: we do not feed it to the environment, but it controls which expert produced the action. In an actor–critic setup, a standard policy-gradient objective against this mixture is

$$L_\text{PG}(\theta,\phi) = -\mathbb{E}_{t}\left[\hat{A}_t\, \log \bar{\pi}_{\theta,\phi}(a_t \mid s_t)\right],$$

where $\hat{A}_t$ is an advantage estimate. Substituting the mixture form gives

$$\log \bar{\pi}_{\theta,\phi}(a_t \mid s_t) = \log \sum_{z_t} \rho_\phi(z_t \mid s_t)\,\pi_{\theta}(a_t \mid s_t, z_t).$$

Conceptually, this means:

• We are maximizing a marginal likelihood over actions in which the expert index $z_t$ is unobserved.
• Gradients must flow through a log-sum over experts, which is expensive and can be unstable when there are many experts or when routing is discrete.

This is exactly the same structure as maximum likelihood in a mixture model with unknown component assignments, and it is this marginal form that will give rise to a nested estimator when approximated.

1.1 Top-$k$ routing as a subset distribution

Many practical MoE architectures route not a single expert, but a subset of $k$ experts at each timestep. In that case the latent routing variable is a subset $S_t \subset \{1,\dots,K\}$ with $|S_t| = k$, rather than a single index. Given router scores $\{s_i(s_t)\}_{i=1}^K$, a common abstraction is to view the induced distribution over orderings as a Plackett–Luce model:

$$p(\sigma \mid s_t) \propto \prod_{j=1}^K \exp\big(s_{\sigma(j)}(s_t)\big),$$

where $\sigma$ ranges over permutations of $\{1,\dots,K\}$. Selecting the “top-$k$” experts then corresponds to taking the first $k$ positions $\sigma(1{:}k)$ (possibly with additional deterministic or stochastic pruning).

In real systems, the router scores themselves are often subject to small amounts of implementation-level noise (for example, due to nondeterministic GPU kernels or reduction orderings), which effectively induces an implicit sampling distribution over experts even when the router is nominally deterministic.² From the latent-variable perspective, this means that the subset $S_t$ is genuinely random under a distribution parameterized by the router, and routing replay is logging samples from this random subset distribution.

One concrete way to sample from this subset distribution is via the Gumbel–top-$k$ trick. Define nonnegative weights $w_i = \exp(s_i(s_t))$ and draw i.i.d. Gumbel noise $g_i \sim \text{Gumbel}(0,1)$. Form perturbed scores $\tilde{s}_i = \log w_i + g_i$ and sort the experts by $\tilde{s}_i$ in descending order. It can be shown that the resulting random permutation is exactly distributed according to the Plackett–Luce model with parameters $\{w_i\}$, so taking the first $k$ experts in this sorted order yields a size-$k$ subset whose distribution matches the Plackett–Luce–induced “top-$k$ without replacement” distribution.

---

2. Routing replay as complete-data training

We now relate this formulation to the practical routing replay technique used in RL. In routing replay, you store the router’s decision $z_t$ at data collection time:

$$\mathcal{D} = \{(s_t, a_t, r_t, s_{t+1}, z_t)\}_{t}.$$

More concretely:

• During data generation, $z_t$ is sampled from the behavior router $\rho_{\phi_\text{behav}}(z_t \mid s_t)$.
• During training, you don’t infer $z_t$ from $(s_t, a_t)$; instead, you treat $z_t$ as observed data.

In latent-variable language, routing replay corresponds to choosing a degenerate inference model:

$$q_\psi(z_t \mid s_t, a_t) = \delta(z_t - \tilde{z}_t),$$

where $\tilde{z}_t$ is the stored routing choice in the replay buffer.

Then the inner expectation collapses:

$$\mathbb{E}_{q_\psi}[\cdot] = \log \rho_\phi(\tilde{z}_t \mid s_t) + \log \pi_\theta(a_t \mid s_t, \tilde{z}_t) - \log 1,$$

because the entropy term $-\mathbb{E}_{q_\psi}[\log q_\psi]$ is zero for a delta.

Plugging this into the ELBO-PG objective gives the routing replay objective:

$$L_{\text{RR}}(\theta,\phi) = -\mathbb{E}_{(s_t,a_t,\tilde{z}_t) \sim \mathcal{D}} \left[ \hat{A}_t\, \big( \log \rho_\phi(\tilde{z}_t \mid s_t) + \log \pi_\theta(a_t \mid s_t, \tilde{z}_t) \big) \right].$$

This is exactly the complete-data objective you would write down for a mixture model where the component assignment $z_t$ (the expert index) is known.

---

3. No routing replay and the nested estimator

We can now address the central question: what goes wrong if we do not perform routing replay? From the latent-variable perspective, discarding the sampled routes $z_t$ and attempting to learn only from the marginal mixture

$$\log \bar{\pi}_{\theta,\phi}(a_t \mid s_t) = \log \sum_{z_t} \rho_\phi(z_t \mid s_t)\,\pi_{\theta}(a_t \mid s_t, z_t)$$

forces us to optimize

$$-\hat{A}_t \log \sum_z \rho_\phi(z \mid s_t)\pi_\theta(a_t \mid s_t,z),$$

i.e., a marginal mixture likelihood with a discrete latent and a log-sum over experts.

This has several issues:

• Credit assignment to the router becomes poorly conditioned: every expert contributes inside the log-sum, even if only one expert actually produced the action.
• In realistic MoE architectures, the sum over experts is expensive or intractable, so it is natural to approximate it by Monte Carlo over routes; as we explain below, this introduces a nested Monte Carlo estimator with poor bias–variance behaviour (in the sense of Rainforth et al., 2018).
• The router can exploit the log-sum-exp structure in undesirable ways (for example, collapsing to a few experts or inducing brittle routing changes), because its gradients are driven by a “soft” mixture rather than the actual routing decisions used to act.

From a sequence-level RL perspective, the nesting can be seen explicitly. Consider an episodic setting with trajectories $\tau = (s_{1:T}, a_{1:T})$ and a single return $R(\tau)$ shared across all timesteps. If we keep the routes (routing replay), a complete-data policy-gradient estimator has the form

$$\nabla J_\text{cd} = \mathbb{E}_{\tau,z_{1:T}} \left[ R(\tau) \sum_{t=1}^T \nabla \log \rho_\phi(z_t \mid s_t) + \nabla \log \pi_\theta(a_t \mid s_t,z_t) \right],$$

which is a single expectation over $(\tau,z_{1:T})$ and therefore non-nested. If we discard $z_{1:T}$ and instead work with the marginal MoE policy, the exact marginal gradient can be written as

$$\nabla J_\text{marg} = \mathbb{E}_{\tau} \left[ R(\tau) \sum_{t=1}^T \nabla \log \sum_{z_t} \rho_\phi(z_t \mid s_t)\pi_\theta(a_t \mid s_t,z_t) \right].$$

Approximating each inner sum over $z_t$ by Monte Carlo leads to an estimator of the form

$$\hat{\nabla J} = \frac{1}{N}\sum_{n=1}^N R(\tau^{(n)}) \sum_{t=1}^T \nabla \log \hat{\mu}_t(\tau^{(n)}),$$

where, for each trajectory $\tau^{(n)} = (s_{1:T}^{(n)},a_{1:T}^{(n)})$ and timestep $t$,

$$\hat{\mu}_t(\tau^{(n)}) = \frac{1}{M_t}\sum_{m=1}^{M_t} \rho_\phi(z_{t,m}^{(n)} \mid s_t^{(n)}) \pi_\theta(a_t^{(n)} \mid s_t^{(n)}, z_{t,m}^{(n)}), \quad z_{t,m}^{(n)} \sim \rho_\phi(\cdot \mid s_t^{(n)}).$$

This matches the generic nested structure

$$\mathbb{E}_{x}\big[F(\mu(x))\big],$$

with $x$ identified with trajectories $\tau$, and with the inner objects $\{\mu_t(\tau)\}_{t=1}^T$ defined by expectations over routes at each timestep. There is a single outer expectation over trajectories and, for each timestep $t$, an inner expectation over route choices $z_t$, all coupled through the shared return $R(\tau)$.

At a more abstract level, a nested Monte Carlo estimator arises whenever we wish to estimate a quantity of the form

$$\mathbb{E}_{x}\big[F(\mu(x))\big], \qquad \mu(x) := \mathbb{E}_{y\mid x}[g(x,y)],$$

and we approximate the inner expectation $\mu(x)$ by Monte Carlo:

$$\hat{\mu}_M(x) = \frac{1}{M}\sum_{m=1}^M g(x,y^{(m)}), \quad y^{(m)} \sim p(y \mid x),$$

before applying the outer nonlinearity $F(\cdot)$. The resulting estimator

$$\hat{I}_{N,M} = \frac{1}{N} \sum_{n=1}^N F\big(\hat{\mu}_M(x^{(n)})\big), \quad x^{(n)} \sim p(x),$$

is nested: it contains an inner Monte Carlo loop inside an outer one. From the latent-variable perspective, passing from the complete-data objective $J_\text{cd} = \mathbb{E}_{x,z}[G(x,z)]$ to its marginalized form $J_\text{marg} = \mathbb{E}_x[\mathbb{E}_{z\mid x}[G(x,z)]]$ and then approximating the inner expectation $\mathbb{E}_{z\mid x}$ by Monte Carlo is precisely what introduces this nested structure. In this sense, marginalizing $z$ from $p(x,z)$ to $p(x)$ when we still have access to samples of $z$ is a weakening of the estimator: it trades a single-level objective on complete data for a nested objective on partial data.

In contrast, routing replay keeps the complete-data objective from Section 2:

$$-\hat{A}_t \big( \log \rho_\phi(z_t \mid s_t) + \log \pi_\theta(a_t \mid s_t,z_t) \big),$$

which has several advantages:

• Clean credit assignment: the router is trained on the actual route that was taken when the action was executed and the return was observed.
• A well-behaved gradient estimator: there is no need to differentiate through a log-sum over experts; only the log-probabilities of a single expert and its route are involved.
• Conceptual alignment with standard latent-variable training: the setting corresponds to an “E-step already done” regime in which component assignments are observed.

In summary, routing replay is the MoE RL analogue of training on complete data instead of marginalizing over unknown latents. If one discards the router decisions and trains only on the marginal mixture, one implicitly tackles a more difficult latent-variable optimization problem whose natural approximations are nested estimators with poor statistical properties.

This perspective helps explain why, in practice, MoE RL setups that do not log and replay routes often behave poorly, whereas routing replay tends to stabilize training: it replaces a nested, ill-conditioned latent-variable objective with a tractable complete-data one.

4. Quantitative analysis of the nested estimator

We now quantify how much worse nested estimators are compared to non-nested estimators. Rainforth et al. (2018)¹ show that, for fixed total budget $T = N M$, nested Monte Carlo estimators of the form above are generally biased and exhibit worse mean-squared-error scaling than non-nested estimators, because the nonlinearity $F$ amplifies inner Monte Carlo error.

It is useful to contrast this with standard (non-nested) Monte Carlo. With a single expectation $\mathbb{E}[h(X)]$ and total budget $T$ samples, the mean-squared error (MSE) of the usual Monte Carlo estimator scales as

$$\mathrm{MSE}_\text{non-nested} = O(T^{-1}).$$

In the nested setting, even with an optimal choice of inner and outer sample sizes $(M,N)$ subject to $T = N M$, Rainforth et al. show that

$$\mathrm{MSE}_\text{nested} = O(T^{-2/3}).$$

Equivalently, to reach a target accuracy $\varepsilon$, a non-nested estimator requires $T = O(\varepsilon^{-2})$, whereas a nested estimator requires $T = O(\varepsilon^{-3})$. This gap quantifies how much “weakening” the estimator by introducing an inner Monte Carlo loop degrades its statistical efficiency.

In the MoE RL setting, there are two natural granularities for $x$. At the per-decision level, $x$ can be taken as $(s_t,a_t)$ and $y$ as the expert index $z_t$, with $F$ given by the composition with the logarithm and policy-gradient weighting. At the trajectory level, $x$ corresponds to an entire trajectory $\tau = (s_{1:T},a_{1:T},r_{1:T})$ and $y$ to the sequence of routing decisions $z_{1:T}$. Crucially, in an MoE policy each $z_t$ is sampled before the action $a_t$ and affects all subsequent states via the dynamics $s_{t+1} \sim p(\cdot \mid s_t,a_t)$. The inner expectation is therefore over an entire routing sequence $z_{1:T}$ whose dimensionality and influence both grow with horizon $T$, so naive nested Monte Carlo here is genuinely “nested over time” and becomes particularly pathological in long-horizon MoE RL.

For fixed $(s_t,a_t)$, define

$$h(z_t; s_t,a_t) := \nabla_\phi \log \rho_\phi(z_t \mid s_t) + \nabla_\theta \log \pi_\theta(a_t \mid s_t, z_t).$$

Then the per-sample gradients for the marginalized and routing-replay objectives can be written as

$$g_\text{marg}(s_t,a_t) = \hat{A}_t\,\mathbb{E}_{z_t \mid s_t,a_t}[h(z_t; s_t,a_t)],$$ $$g_\text{RR}(s_t,a_t,z_t) = \hat{A}_t\,h(z_t; s_t,a_t).$$

By the law of total variance, conditioning on $(s_t,a_t)$,

$$\mathrm{Var}[g_\text{RR}] = \mathrm{Var}[g_\text{marg}] + \mathbb{E}_{(s_t,a_t)}\Big[ \hat{A}_t^2\, \mathrm{Var}_{z_t \mid s_t,a_t}\big[h(z_t; s_t,a_t)\big] \Big].$$

The second term is non-negative, so in the idealized regime where exact marginalization over experts is tractable, the marginal objective yields a lower-variance gradient estimator than the complete-data objective. In practice, however, the exact sum over experts is often prohibitively expensive in large MoE architectures.

When the sum over experts is approximated by Monte Carlo, the marginal objective induces a nested Monte Carlo estimator. Writing

$$\mu(s_t,a_t) := \mathbb{E}_{z_t \mid s_t}[\pi_\theta(a_t \mid s_t,z_t)], \qquad F(\mu) := -\hat{A}_t \log \mu,$$

one replaces $\mu$ by a finite-sample estimate

$$\hat{\mu}_M(s_t,a_t) = \frac{1}{M} \sum_{m=1}^M \pi_\theta(a_t \mid s_t, z_t^{(m)}),$$

and uses $F(\hat{\mu}_M)$ inside an outer expectation over the replay distribution. This is exactly the nested structure studied by Rainforth et al. (2018): a nonlinear transformation of an inner Monte Carlo estimate inside an outer expectation.

The key implication is that, for a fixed total sample budget, such nested estimators exhibit worse bias and mean-squared-error scaling than non-nested estimators. Thus approximate marginalization over experts can be substantially less statistically efficient than either: (i) exact marginalization (when feasible), or (ii) complete-data training with routing replay, which remains non-nested.

From this viewpoint, the primary benefit of routing replay is that it avoids nested Monte Carlo altogether by converting the latent route into observed data and optimizing a single-level objective.

4.1 Case study: variance behavior for MoE RL

It is also useful to compare, at a qualitative level, how the variance of different MoE RL gradient estimators behaves as a function of the horizon $T$ and the choice of estimator. Consider an episodic setting with a single return $R(\tau)$ per trajectory and define, for simplicity, the per-timestep complete-data contribution

$$g_{\text{RR},t}(\tau) := R(\tau)\,\big(\nabla_\phi \log \rho_\phi(z_t \mid s_t) + \nabla_\theta \log \pi_\theta(a_t \mid s_t,z_t)\big),$$

so that the routing-replay gradient estimator is

$$g_\text{RR}(\tau) = \sum_{t=1}^T g_{\text{RR},t}(\tau).$$

• Case 1: Complete-data (routing replay), single sample per trajectory.
  Even without nesting, variance typically grows at least linearly with $T$. Writing

  $$\mathrm{Var}[g_\text{RR}] = \sum_{t=1}^T \mathrm{Var}[g_{\text{RR},t}] + 2\sum_{1\le t<u\le T} \mathrm{Cov}(g_{\text{RR},t},g_{\text{RR},u}),$$

  and assuming each $g_{\text{RR},t}$ has non-zero variance and covariances are not strongly negative, we obtain $\mathrm{Var}[g_\text{RR}] \gtrsim c T$ for some $c>0$. This is the usual “variance grows with horizon” behavior of sequence-level REINFORCE.

• Case 2: Exact marginalization over experts (intractable ideal).
  If we could compute the sum over experts exactly, the marginal gradient estimator would have lower variance than the complete-data one by the law of total variance: integrating out $z_t$ removes conditional variability. However, this regime is not computationally accessible in realistic MoE architectures.

• Case 3: Single-sample plug-in approximation to the marginal.
  In practice, approximating $\log \sum_{z_t} \rho_\phi(z_t\mid s_t)\pi_\theta(a_t\mid s_t,z_t)$ with a single-sample plug-in of the form $\log \hat{\mu}_1(s_t,a_t)$ (or its gradient) introduces additional noise: each per-timestep term becomes a nonlinear function of a single random route sample. Summing these noisy terms over $t$ and scaling by a shared $R(\tau)$ again gives at least linear growth in variance with $T$, but with larger per-timestep variance and bias compared to the complete-data case.

Taken together, these cases suggest the following picture: routing replay does not eliminate the usual horizon-related variance issues of RL, but it avoids the additional variance and bias introduced by approximating a high-dimensional marginal log-sum over experts with noisy single-sample plug-ins. The nested Monte Carlo viewpoint clarifies why the marginal training problem is statistically harder than the complete-data one, even when both ultimately rely on single-sample estimators in practice.

---

6. One-line summary

Routing replay turns the expert index $z_t$ from an unobserved latent variable in the mixture policy $\bar{\pi}(a_t\mid s_t)$ into an observed variable stored in replay, so the MoE RL training objective becomes a weighted complete-data (latent-variable) objective:

$$\max_{\theta,\phi} \;\mathbb{E}\left[\hat{A}_t \left( \log \rho_\phi(z_t \mid s_t) + \log \pi_\theta(a_t \mid s_t,z_t) \right)\right],$$

which is precisely the standard latent-variable formulation with known component assignments.

Footnotes

1. Tom Rainforth, Robert Cornish, Hongseok Yang, Andrew Warrington, and Frank Wood. On Nested Monte Carlo Estimation. ICML 2018. arXiv:1606.00323. ↩ ↩²

2. Horace He. “Defeating Nondeterminism.” Thinking Machines blog. ↩