Tempered Sequential Monte Carlo for Trajectory and Policy Optimization with Differentiable Dynamics

Abstract: We propose a sampling-based framework for finite-horizon trajectory and policy optimization under differentiable dynamics by casting controller design as inference. Specifically, we minimize a KL-regularized expected trajectory cost, which yields an optimal "Boltzmann-tilted" distribution over controller parameters that concentrates on low-cost solutions as temperature decreases. To sample efficiently from this sharp, potentially multimodal target, we introduce tempered sequential Monte Carlo (TSMC): an annealing scheme that adaptively reweights and resamples particles along a tempering path from a prior to the target distribution, while using Hamiltonian Monte Carlo rejuvenation to maintain diversity and exploit exact gradients obtained by differentiating through trajectory rollouts. For policy optimization, we extend TSMC via (i) a deterministic empirical approximation of the initial-state distribution and (ii) an extended-space construction that treats rollout randomness as auxiliary variables. Experiments across trajectory- and policy-optimization benchmarks show that TSMC is broadly applicable and compares favorably to state-of-the-art baselines.

• The paper's main contribution is the novel TSMC algorithm that reframes trajectory and policy optimization as an inference problem using a KL-regularized cost.
• It employs an adaptive annealing schedule and Hamiltonian Monte Carlo rejuvenation to efficiently sample from sharply multimodal, Boltzmann-tilted distributions.
• Numerical experiments demonstrate TSMC’s superiority over traditional methods by achieving lower trajectory costs and reliable policy performance in nonsmooth, contact-rich domains.

Tempered Sequential Monte Carlo for Trajectory and Policy Optimization with Differentiable Dynamics

Problem Formulation and Motivation

The paper "Tempered Sequential Monte Carlo for Trajectory and Policy Optimization with Differentiable Dynamics" (2604.21456) addresses finite-horizon trajectory optimization (TO) and policy optimization (PO) by leveraging the differentiability of dynamical systems and controllers. The classical optimization objective $$\min_{\theta} \mathbb{E}_{x_0 \sim \mu} [J(\tau(x_0, \theta))]$$ is reframed as inference via a KL-regularized expected trajectory cost. This yields an optimal distribution over controller parameters proportional to a Boltzmann-tilted density, concentrating around low-cost solutions as the temperature parameter $\lambda$ decreases.

Strong motivation is provided by automatic differentiation capabilities in modern simulator frameworks and the limitations imposed by gradient-based optimization in nonconvex and nonsmooth settings. The core technical challenge is efficient sampling from sharp, multimodal distributions induced by Boltzmann tilting—an area where standard importance sampling and MCMC approaches suffer from high cost and poor mixing, especially as temperature approaches zero.

Tempered Sequential Monte Carlo: Methodology

Tempered Sequential Monte Carlo (TSMC) is introduced as a principled algorithmic framework to sample from the Boltzmann-tilted distribution. TSMC achieves this via an annealing schedule that interpolates between an easy-to-sample prior and the target using a set of tempering parameters $(\beta_0, ..., \beta_K)$, adaptive ESS-based selection of tempering levels, and importance reweighting. Particle degeneracy is mitigated by resampling, and diversity is restored through Hamiltonian Monte Carlo (HMC) rejuvenation kernels exploiting exact gradients from differentiable simulators.

Figure 1: Tempered Sequential Monte Carlo samples from the Shekel-tilted distribution, illustrating multimodality and the effect of annealing and HMC rejuvenation.

The paper details the connection between the optimal density $p^{\star}(\theta)$ and classical Gibbs posteriors in statistical physics and the "control-as-inference" paradigm. For trajectory optimization, gradients w.r.t. controller parameters are computed by reverse-mode autodiff or adjoint recursion exploiting system structure (see Theorem~\ref{thm:gradient-computation-trajectory-optimization}). For policy optimization, energy and gradients are approximated by Monte Carlo rollouts, and two variants are proposed: deterministic empirical energy approximation and extended-space TSMC with auxiliary variables encoding rollout randomness.

Numerical Experiments

Comprehensive experiments are conducted for both trajectory and policy optimization across classical control benchmarks and nonsmooth contact-rich domains.

Figure 2: Trajectory optimization results: boxplots of final-particle costs for TSMC and baselines; best-cost trajectories for Acrobot and pusher–slider.

On trajectory optimization, TSMC consistently outperforms baselines (parallel NUTS, MALA, MPPI, IPOPT) in difficult settings (Acrobot, pusher-slider), reliably attaining lower trajectory costs and successful goal attainment. Notably, TSMC is the only method to achieve upright swing-up on Acrobot and goal pose on pusher-slider, underlining robust convergence across sharp energy landscapes.

Figure 3: Energy distributions across tempering steps of TSMC for TO (inverted pendulum, acrobot, pusher–slider), demonstrating improved concentration.

For policy optimization, TSMC competes favorably with high-performing RL algorithms such as SAC and PPO. On sparse-reward inverted pendulum, both deterministic and extended-space variants yield policies matching SAC performance. For Acrobot and double pendulum, TSMC finds policies with lower costs compared to PPO and SAC, even as gradient-based RL struggles to find feasible solutions within practical computation budgets.

Figure 4: Policy optimization results: RL returns for PPO and SAC, TSMC energy distributions, and policy rollouts across tasks.

Long-horizon RL experiments for Acrobot highlight the difficulty of obtaining successful policies with standard methods even with extended training. TSMC's ability to operate on sharply nonconvex objectives is reinforced by empirical evidence, particularly under single-state initial distributions where energy landscape sharpness is most pronounced.

Practical and Theoretical Implications

The work bridges trajectory and policy optimization, offering a unified framework capable of integrating sampling-based exploration and gradient-based refinement under differentiable dynamics. TSMC expands the operation of classical optimization by facilitating robust exploration in low-temperature, multimodal regimes. It retains natural parallelizability but introduces higher computational and memory demands, especially as particle count and rollout lengths increase.

The theoretical guarantees (consistency under increasing particle count) are classical, but the practical efficacy stems from effective annealing, adaptive tempering, and HMC rejuvenation. The framework adapts seamlessly to stochastic dynamics and risk-sensitive objectives via extended-space constructions, laying a foundation for robust sampling under uncertainty.

There is strong evidence that TSMC achieves consistent improvement over existing trajectory and policy optimization methods, and, in particular, yields significant improvements in nonsmooth and contact-rich settings where gradients are unreliable or uninformative. The deterministic approximation and extended-space variants are practical for policy optimization even under intractable expectation computations, albeit at increased computational cost.

Implications for Future Developments

TSMC's scalability and robust optimization capabilities suggest clear paths for future integration with differentiable contact solvers and high-dimensional RL benchmarks. Adaptation to neural samplers and amortized MCMC regimes may further accelerate computation and generalization. Hardware acceleration and custom GPU kernels could alleviate computational bottlenecks enabling application to large-scale, real-world robotics.

On the theory side, extensions to certified global optimization via convex relaxations or convergence in risk-sensitive settings could deepen the framework's utility. The paradigm shift towards optimization as inference, instantiated via TSMC, is instrumental in data-driven control and policy search, particularly as differentiable simulation becomes ubiquitous.

Conclusion

The paper introduces TSMC as a robust, principled approach for trajectory and policy optimization leveraging the structure and differentiability of modern simulators. It addresses the crucial challenge of efficiently sampling from sharply multimodal distributions, outperforming existing baselines in both classical and contact-rich domains. While computationally demanding, its practical impact is substantial in settings where gradient information is unreliable or nonconvexity dominates. TSMC positions itself as an essential tool for future developments in AI and robotics optimization, especially as model and simulation fidelity continue to improve.