Technical Summary: PISTO – Proximal Inference for Stochastic Trajectory Optimization

Problem Statement

Stochastic trajectory optimization methods, such as STOMP (Stochastic Trajectory Optimization for Motion Planning), offer significant flexibility over gradient-based approaches by handling non-differentiable and discontinuous cost functions (e.g., collision penalties). However, the theoretical underpinnings of STOMP's update rules have remained somewhat opaque, and the algorithm can suffer from unstable optimization dynamics. While STOMP effectively smooths cost landscapes through sampling, there is a need to understand its implicit objective and to introduce mechanisms that stabilize the iterative updates, particularly in complex, non-convex environments.

Methodology

Theoretical Foundation: STOMP as Variational Inference

The authors first establish a theoretical connection between STOMP and Variational Inference (VI). They demonstrate that STOMP implicitly minimizes the reverse Kullback-Leibler (KL) divergence between a proposal Gaussian distribution, $N(Y, \Sigma)$, and a target Boltzmann trajectory distribution, $Y^\star \propto \exp(-S(\tilde{Y}))N(0, R^{-1})$. Here, $S(\tilde{Y})$ represents the accumulated state-dependent cost (e.g., obstacle avoidance), and $N(0, R^{-1})$ acts as a Gaussian smoothness prior. This reformulation reveals that STOMP's cost-weighted averaging is equivalent to a moment-matching solution within a VI framework.

The PISTO Algorithm

Building on this VI interpretation, the paper proposes Proximal Inference for Stochastic Trajectory Optimization (PISTO). The core innovation is the augmentation of the VI objective with a proximal regularization term.

1. Proximal Formulation: Instead of directly minimizing the KL divergence to the target, PISTO solves a proximal update problem:
   $$Y_{k+1} = \arg \min_Y \left[ D_{KL}(N(Y, \Sigma) \parallel Y^\star) + \frac{1}{\eta} D_{KL}(N(Y, \Sigma) \parallel N(Y_k, \Sigma)) \right]$$
   where $\eta$ is a step-size parameter. This formulation admits a trust-region interpretation, penalizing large deviations from the current iterate $Y_k$ to ensure stable convergence.

2. Surrogate Distribution: The proximal update is shown to be equivalent to minimizing the KL divergence to a surrogate distribution $Y^\star_k$. This surrogate is a geometric interpolation between the true Boltzmann target and the current Gaussian approximation.

3. Closed-Form Updates via Importance Sampling: The solution to the proximal problem is a moment-matching update where the new mean $Y_{k+1}$ is the expectation of the surrogate distribution. Since the surrogate is generally intractable to sample from directly, PISTO employs importance-weighted Monte Carlo sampling.

   • Samples are drawn from a Gaussian proposal centered at the current iterate.
   • Weights are computed based on the cost function and the proximal term.
   • The update is a weighted average of these samples.

4. Derivative-Free Nature: Crucially, the algorithm requires only the evaluation of the cost function $S(\cdot)$, making it derivative-free. It inherits STOMP's ability to handle non-differentiable and discontinuous costs without modification.

5. Extensions:

   • Contact-Rich Tasks: The framework is extended to policy optimization in action space for systems with complex dynamics (e.g., MuJoCo), where the state trajectory is determined by rolling out dynamics from control sequences.
   • Annealing: The method utilizes adaptive covariance scheduling (cosine annealing) and a proximal step-size annealing scheme to balance exploration and exploitation. An elite-set selection mechanism is also employed to reduce variance by focusing on low-cost samples.

Key Contributions

1. Theoretical Insight: The paper reveals that STOMP implicitly minimizes the reverse KL divergence from a Boltzmann trajectory distribution, establishing a clear Variational Inference structure for the algorithm.
2. PISTO Algorithm: The authors propose PISTO, which introduces a proximal penalty (KL regularization) between successive Gaussian proposals. This provides a trust-region interpretation that stabilizes optimization dynamics.
3. Closed-Form Derivative-Free Updates: The proximal formulation yields closed-form mean updates computable via importance-weighted Monte Carlo sampling. This results in a simple, derivative-free algorithm that maintains flexibility while offering principled stability guarantees.

Experimental Results

Motion Planning Benchmarks

PISTO was evaluated on robot arm motion planning benchmarks using the MoveIt framework across seven manipulation environments (e.g., Kitchen, Bookshelf, Box).

• Success Rate: PISTO achieved an 89% success rate, significantly outperforming CHOMP (63%), STOMP (68%), and Natural Gradient Descent (76%).
• Efficiency: PISTO ran approximately twice as fast as competing stochastic methods while producing shorter and smoother paths.
• Metrics: It demonstrated superior performance in path clearance and planning time compared to baselines.

Contact-Rich Locomotion and Manipulation

The algorithm was tested on MuJoCo tasks involving complex contact dynamics (e.g., PushT, Humanoid Running, Humanoid Standing Up).

• Reward Performance: PISTO consistently outperformed Cross-Entropy Method (CEM) and Model Predictive Path Integral (MPPI) baselines in reward across all tested tasks. Notably, it converted the negative reward of MPPI on the PushT task (−0.17) into a positive reward (0.46).
• Speed: PISTO demonstrated significant speedups, running 1.4× to 3.2× faster than baselines on most tasks (e.g., 36.5s vs. ~115s for Humanoid Running).
• Robustness: The method successfully solved high-dimensional (17-DOF) humanoid tasks within minutes, validating its robustness in hybrid dynamical systems with discontinuous contact.

Significance and Claims

The paper claims that PISTO represents a principled advancement in stochastic trajectory optimization by unifying the flexibility of sampling-based methods with the stability of proximal inference. By revealing the variational structure of STOMP, the authors provide a theoretical basis for algorithmic improvements. The resulting algorithm is highlighted for being:

• Simple and Derivative-Free: Capable of handling non-differentiable costs without modification.
• Stable: The proximal formulation prevents overly aggressive steps, ensuring reliable convergence in non-convex landscapes.
• Efficient: It achieves state-of-the-art success rates and speeds on both standard motion planning benchmarks and challenging contact-rich control tasks.

The authors conclude that PISTO offers a robust solution for trajectory optimization in complex robotic systems, with future work suggested for incorporating gradients where available and extending the method to receding-horizon Model Predictive Control (MPC).