Technical Summary: Trajectory-Regularized Stochastic Optimal Control via KL Divergence

1. Problem Formulation

The paper addresses a limitation in standard Stochastic Optimal Control (SOC), which focuses solely on minimizing expected cumulative cost under uncertainty without constraining how the resulting system dynamics deviate from a desired reference behavior. In many applications, reference behaviors (e.g., passive dynamics, baseline controllers, or distributions derived from offline data) encode valuable structural priors. Large deviations from these references can lead to brittle policies.

The authors propose Trajectory-Regularized Stochastic Optimal Control (TRSOC). The objective is to minimize a cost function that balances task performance with the proximity of the controlled trajectory distribution to a reference trajectory distribution. This proximity is measured using the Kullback–Leibler (KL) divergence between the probability laws of the controlled process ($P_u$) and the reference process ($P_0$).

The system is modeled as a controlled diffusion process:
$$dX_t = f(X_t, u_t)dt + \sigma(X_t)dW_t$$
where the control $u_t$ enters through the drift. The reference dynamics share the same diffusion matrix $\sigma$ but have a different drift $f_0$.

2. Methodology

A. Reduction of Trajectory KL to Local Cost

A central theoretical contribution is the reduction of the trajectory-level KL divergence to a local, tractable running cost.

• Girsanov's Theorem: Under the assumption that the diffusion matrix $\sigma(x)$ is invertible and the drift mismatch satisfies Novikov's condition, the Radon–Nikodym derivative between the controlled and reference measures exists.
• Lemma 1: The KL divergence $KL(P_u \| P_0)$ over a finite horizon $[0, T]$ is shown to be equal to the expected integral of a quadratic drift-mismatch penalty:
  $$KL(P_u \| P_0) = \frac{1}{2} \mathbb{E}_{P_u} \left[ \int_0^T \|\phi_t^u\|^2 dt \right]$$
  where $\phi_t^u = \sigma(X_t)^{-1}(f(X_t, u_t) - f_0(X_t))$.
• Augmented Running Cost: This result allows the trajectory regularization to be incorporated directly into the standard SOC framework as an additive term to the running cost:
  $$L(x, u) = \ell(x, u) + \frac{\lambda}{2} \|\sigma(x)^{-1}(f(x, u) - f_0(x))\|^2$$
  Here, $\lambda \geq 0$ is the regularization parameter.

B. Dynamic Programming and HJB Structure

Because the KL penalty reduces to an integral of a running cost, the TRSOC problem preserves the Markov property and the Dynamic Programming (DP) principle.

• The value function $V(t, x)$ satisfies the standard Hamilton–Jacobi–Bellman (HJB) equation:
  $$-\partial_t V(t, x) = \inf_{u \in U} \{ L(x, u) + \mathcal{L}_u V(t, x) \}$$
  where $\mathcal{L}_u$ is the infinitesimal generator of the diffusion.
• Unlike "KL control" approaches that optimize directly over trajectory measures (often leading to risk-sensitive objectives and exponential tilting), TRSOC remains a standard SOC problem with an augmented cost, retaining a linear expectation and a standard HJB characterization.

C. Optimal Control Structure (Control-Affine Case)

For systems with control-affine drift $f(x, u) = f_0(x) + B(x)u$ and quadratic control costs, the authors derive an explicit optimal feedback law.

• The KL term acts as an additional control penalty, modifying the effective control weight matrix:
  $$\bar{R}(x) = R(x) + \lambda B(x)^\top (\sigma(x)\sigma(x)^\top)^{-1} B(x)$$
• The optimal control is given by:
  $$u^*(x) = -\bar{R}(x)^{-1} B(x)^\top \nabla V(x)$$
• This formulation implies that controls inducing large drift changes in low-noise directions are penalized more heavily.

D. Linear-Quadratic (LQ) Specialization

In the Linear-Quadratic setting, the problem reduces to a discounted Linear Quadratic Regulator (LQR) with a modified control weight $\bar{R}$. The solution involves solving a modified algebraic Riccati equation (ARE), where the regularization parameter $\lambda$ directly influences the feedback gain, leading to more conservative control as $\lambda$ increases.

3. Key Contributions

1. TRSOC Formulation: Introduction of a behavior-regularized SOC framework that incorporates prior behaviors (nominal controllers or data-driven baselines) via trajectory-level KL divergence.
2. Local Cost Reduction: Demonstration that under drift-only control with non-degenerate diffusion, the trajectory KL divergence reduces to a quadratic drift-mismatch penalty. This preserves the DP structure and allows for standard HJB characterization.
3. Optimal Policy Characterization: Derivation of explicit optimal feedback laws for control-affine systems and analysis of the LQ case, showing how regularization modifies the control cost and stability properties.
4. Stability Analysis: Proof that the trajectory regularization can improve the dissipativity of the closed-loop drift, facilitating Lyapunov inequalities and ensuring mean-square boundedness of the state.

4. Results and Experiments

The paper validates the theoretical findings through numerical experiments on a stochastic double-integrator system tracking a figure-eight trajectory.

• Performance–Deviation Trade-off: Experiments demonstrate that the regularization parameter $\lambda$ induces a continuous trade-off.
  • As $\lambda \to 0$, the controller behaves like a standard performance-driven SOC, achieving low task cost but potentially deviating significantly from the reference.
  • As $\lambda \to \infty$, the controller forces the induced drift to align with the reference dynamics, minimizing deviation at the cost of task performance.
• Data-Driven Reference: The framework is successfully applied to a setting where the reference dynamics are learned from offline data (using a neural network to approximate the reference acceleration). The results show that TRSOC effectively balances task optimization with adherence to the learned reference behavior.
• Trajectory Geometry: Visualizations of trajectories reveal that the regularization acts on the induced dynamics geometry. Regularized controllers produce trajectories that lie between the aggressive corrections of the unregularized optimal controller and the nominal reference behavior, particularly in high-curvature regions.

5. Significance and Claims

The authors claim that TRSOC provides a principled method to integrate prior knowledge or offline data into stochastic control without abandoning the standard DP/HJB framework.

• Structural Preservation: Unlike other trajectory regularization methods that may require optimizing over measures or lead to non-standard risk-sensitive objectives, TRSOC maintains the standard SOC structure, making it applicable to general nonlinear dynamics.
• Interpolation: The framework offers a scalarization of a multi-objective optimization problem (task performance vs. reference adherence), allowing practitioners to tune the behavior of the controller via a single parameter $\lambda$.
• Robustness: The method is shown to be robust to approximation errors when the reference dynamics are learned from data, suggesting potential utility in offline-to-online reinforcement learning and imitation learning contexts.

The paper concludes that trajectory regularization acts as a "noise-aware" modification of control energy, penalizing drift changes more heavily in directions where the system is less stochastic, thereby offering a geometrically intuitive mechanism for balancing performance and safety/consistency.