PILCO, PETS, Ensembles, and MBPO › Trajectory sampling with ensembles

PETS, short for Probabilistic Ensembles with Trajectory Sampling, estimates action-sequence values with particles propagated through a probabilistic neural ensemble. The ensemble members are plausible fitted functions and their disagreement is an empirical epistemic proxy; ordinary deep-ensemble members are not formal draws from a Bayesian posterior.

PILCO's one-step GP posterior is analytic under its model assumptions, but its long-term state prediction uses approximate Gaussian moment matching. Its data efficiency cannot be attributed to end-to-end exactness. PETS instead estimates action-sequence returns by finite-sample particle rollout through its ensemble. Increasing particle count can reduce Monte Carlo error in a candidate's value estimate, while model bias, ensemble approximation error, and action-search error remain separate limitations.

At the model level, PETS represents the conditional next-state distribution as a diagonal Gaussian. In the released implementation, the network may receive a preprocessed observation with the action and predicts Gaussian parameters for a task-specific target, with variances parameterized to be strictly positive: the Cartpole configuration uses state increments, while HalfCheetah uses an absolute first next-state coordinate and increments for the others. Deterministic postprocessing reconstructs the next state. The following $\mu_\theta$ and $\Sigma_\theta$ denote moments of that induced next-state distribution, not necessarily raw network outputs:

$$q_\theta(s_{t+1} \mid s_t, a_t) = \mathcal{N}\big(\mu_\theta(s_t, a_t), \Sigma_\theta(s_t, a_t)\big)$$

where:

• $q_\theta(s_{t+1} \mid s_t, a_t)$: the model's predicted distribution over the next state
• $\mu_\theta(s_t, a_t)\in\mathbb{R}^{d_s}$: the predicted mean next-state vector after target postprocessing
• $\Sigma_\theta(s_t, a_t)\in\mathbb{R}^{d_s\times d_s}$: the predicted diagonal next-state covariance after target postprocessing
• $\theta$: the network parameters

The diagonal covariance treats modeled output dimensions as conditionally independent. The released target maps above are coordinatewise affine with unit Jacobian, so the Gaussian data-fit term can also be written as a negative log-likelihood of the induced next states. It penalizes both an incorrect mean and an overconfident variance: an incorrect mean increases squared error relative to predicted variance, and a variance that is too small magnifies that error. The unregularized data-fit term is

$$\mathcal{L}(\theta) = -\sum_{(s, a, s') \in \mathcal{D}} \log \mathcal{N}\big(s' \mid \mu_\theta(s, a), \Sigma_\theta(s, a)\big)$$

where:

• $\mathcal{L}(\theta)$: the Gaussian negative log-likelihood data-fit term as a function of the network parameters
• $\mathcal{D}$: the dataset of observed transitions $(s, a, s')$
• $s'$: the observed next state
• $\mathcal{N}(s' \mid \mu_\theta(s, a), \Sigma_\theta(s, a))$: the likelihood the model assigns to the observed transition

The released PETS training code adds weight decay and penalties on learned log-variance bounds to this Gaussian data-fit term; the full optimization is regularized, not pure maximum likelihood. Training an input-dependent variance output with log-likelihood, rather than fitting only a mean by squared error, permits a heteroscedastic model. With adequate data and a well-fitted conditional model, it can predict higher variance where transitions have greater irreducible variability and lower variance where they are consistent. The learned $\Sigma_\theta$ is then an aleatoric proxy, not a clean measurement: mean-model error can also inflate it, and the training loss does not identify its value in unseen regions. A single network can therefore output a confident, low-variance prediction where it has never seen data. Ensemble disagreement supplies a separate empirical proxy for epistemic uncertainty.

A single probabilistic network with a likelihood data-fit term can learn input-dependent aleatoric variance on observed data, but that term does not directly constrain epistemic uncertainty at unobserved inputs. Extra capacity can produce confident off-support behavior, but it does not necessarily make that behavior worse; the outcome depends on architecture, optimization, regularization, and data.

The deep-ensemble recipe trains $N$ networks with independent randomization and, in PETS, bootstrap resampling. For vector-valued predicted means, define

$$\begin{aligned} \bar{\mu}(x) &= \frac{1}{N}\sum_{i=1}^{N}\mu_{\theta_i}(x), \\ \widehat{\operatorname{Cov}}_{\mathrm{epi}}(x) &=\frac{1}{N}\sum_{i=1}^{N} \big(\mu_{\theta_i}(x)-\bar{\mu}(x)\big) \big(\mu_{\theta_i}(x)-\bar{\mu}(x)\big)^\top. \end{aligned}$$

For a scalar output, the covariance expression reduces to the familiar average squared deviation. For a state-vector output, $\bar{\mu}\in\mathbb{R}^{d_s}$ and the covariance is a $d_s\times d_s$ matrix. The between-member covariance is used as an epistemic proxy. With probabilistic members, a useful approximate total-covariance decomposition is

$$\widehat{\operatorname{Cov}}_{\mathrm{total}}(x) =\frac{1}{N}\sum_{i=1}^{N}\Sigma_{\theta_i}(x) +\widehat{\operatorname{Cov}}_{\mathrm{epi}}(x),$$

which keeps average within-member predictive variance separate from between-member disagreement.

In observed regions, independently trained members can converge toward similar functions; away from the data, different random initializations, training paths, and optional bootstrap resampling can lead to different extrapolations. This behavior is common enough to be useful, but it is not guaranteed or monotone in data density. Correlated members may agree while sharing the same error, and optimization variability can create disagreement even in densely sampled regions.

Ordinary deep-ensemble members are not posterior samples. Bootstrap resampling approximates a frequentist sampling distribution under its own conditions, not generally the Bayesian function posterior conditional on one observed dataset. The method is therefore best described as an empirical ensemble used as an epistemic proxy. A related Randomized Prior Functions construction adds a fixed, randomly drawn prior function to each trainable member's output; it does not merely randomize the final layer and train the remainder. These methods can work well empirically without turning their fitted-member distribution into an analytic or simulation-exact posterior.

PETS evaluates each candidate action sequence with multiple Monte Carlo state particles. Candidates, particles, and ensemble members are different objects. Let $c$ index a CEM candidate, $p\in\{1,\ldots,P\}$ index one of its particles, and $b(p,\tau)$ identify the ensemble member assigned to that particle at simulated time $\tau$. Starting every particle from the measured state, PETS propagates

$$s_{\tau+1}^{(c,p)}\sim \mathcal{N}\!\left( \mu_{\theta_{b(p,\tau)}}\!\left(s_\tau^{(c,p)},a_\tau^{(c)}\right), \Sigma_{\theta_{b(p,\tau)}}\!\left(s_\tau^{(c,p)},a_\tau^{(c)}\right) \right).$$

The candidate score is the particle average of cumulative reward,

$$\widehat{V}\!\left(a_{t:t+H-1}^{(c)}\right) =\frac{1}{P}\sum_{p=1}^{P} \sum_{\tau=t}^{t+H-1} r\!\left(s_\tau^{(c,p)},a_\tau^{(c)},s_{\tau+1}^{(c,p)}\right).$$

Here $r$ is an abstract task-supplied transition score, which may depend on the sampled next state. In the released PETS controller, each predicted next observation receives a task-supplied observation cost and its action receives a task-supplied action cost; maximizing $r$ corresponds to minimizing their sum. The dynamics network does not learn those costs through the Gaussian transition likelihood. This differs from original MBPO's learned dynamics-and-reward ensemble.

Thus CEM's candidate population controls which action sequences are searched, while $P$ controls Monte Carlo uncertainty propagation and the precision of each candidate's return estimate. PETS does not require one trajectory per ensemble member, nor does a particle denote an action-sequence candidate.

The original PETS study defines two bootstrap-assignment semantics:

• TS1 resamples $b(p,\tau)$ at every simulated step. It represents a transition process that marginalizes over or resamples the ensemble hypothesis through time.
• TS$\infty$ samples one member for each particle and holds it fixed across the horizon. It represents a time-invariant sampled dynamics hypothesis and preserves trajectory-level effects of that member.

Both are valid PETS variants. A TS1 trajectory need not coincide with the rollout of any one fixed member, but it is not therefore a statistical fiction. TS1 and TS$\infty$ make different assumptions about temporal dependence in model uncertainty. There is no general identity saying that TS1 must average away long-horizon uncertainty or that TS$\infty$ is always more conservative.

PETS uses the cross-entropy method (CEM) to search over action sequences:

1. Initialize a proposal distribution over sequences.
2. Sample a population of candidate sequences.
3. Estimate each candidate's return with its own state particles.
4. Select a high-return elite set.
5. Refit the proposal to the elites.
6. Repeat the sampling and refitting updates until a stopping criterion is met or the maximum number of iterations is reached. Return a sequence derived from the final proposal; the released PETS optimizer returns its smoothed mean. The surrounding receding-horizon controller executes only the first action.

Candidate and particle evaluations can be batched across neural-network calls. CEM is derivative-free, but it is not immune to search failure. Too few candidates can miss narrow high-value regions; too little proposal variance can collapse the search prematurely; too much variance can spend computation on poor regions; and an extreme elite fraction can make updates either brittle or weak. The appropriate population depends on action dimension, horizon, proposal family, and compute budget rather than a field-wide typical range.

Sampled elite return is one optimization diagnostic, not proof that concentration is beneficial. Proposal variance and returns reevaluated with fresh or additional particles are useful diagnostics for sampling noise and concentration. Distinguishing premature collapse from convergence to a strong solution also requires broader-search or restart evidence.

PETS plans over a finite horizon, executes only the first action, observes the resulting state, and replans. Receding-horizon control therefore limits each open-loop commitment to one action and each uncorrected forecast to the configured planning horizon. It does not numerically bound model error or guarantee global competence. Success can still depend on horizon length, terminal design, feasibility, exploration, and accuracy in regions the task must traverse.

The system loop is:

1. Collect an initial transition dataset.
2. Train the probabilistic ensemble.
3. Use CEM and particle trajectory sampling to rank candidate action sequences.
4. Execute the first action, observe the transition, and append it to the dataset.
5. Retrain the ensemble according to a configured cadence.

The published PETS score is the arithmetic mean of particle returns. That is a risk-neutral Monte Carlo estimate of expected return under the fitted ensemble predictive distribution, not an automatic uncertainty hedge. One low return lowers the mean arithmetically, but two return distributions with the same mean receive the same score even if their spreads differ.

Retraining too rarely leaves the ensemble stale, while retraining too often increases computation. A usable trigger should be observable: for example, retrain after a fixed number of transitions and monitor held-out likelihood, calibration, or prediction drift in recently visited regions. A claim that the dataset has changed the model needs one of those quantities to make it operational.