Source-linked AI summary
Optimal control as a graphical model inference problem
B. Kappen, V. Gomez, M. Opper
TL;DR
The paper addresses the difficulty of computing stochastic optimal controls in large or complex systems. It reformulates a class of controls as KL minimization and graphical-model inference, allowing exact or approximate inference methods. Approximate inference successfully computes controls for block stacking and multi-agent cooperative tasks, while the framework includes path integral control as a special case.
Problem
Stochastic optimal control becomes difficult in high-dimensional systems because the state space can be too large for direct computation.
Method
The paper expresses control as a probability distribution over future trajectories and minimizes a KL-divergence control cost through graphical-model inference.
Results
Approximate inference successfully computes approximate controls for a block stacking task and a multi-agent cooperative task, including settings too complex for exact inference.
Takeaways & Limitations
KL control provides a framework in which exact or approximate inference methods can be applied directly to stochastic control computation and includes path integral control as a special case.
Takeaways & Limitations
The framework covers only a class of stochastic control problems, and the model-free extension requires a sampling-based procedure.
Abstract
from arXiv · showhide
We reformulate a class of non-linear stochastic optimal control problems introduced by Todorov (2007) as a Kullback-Leibler (KL) minimization problem. As a result, the optimal control computation reduces to an inference computation and approximate inference methods can be applied to efficiently compute approximate optimal controls. We show how this KL control theory contains the path integral control method as a special case. We provide an example of a block stacking task and a multi-agent cooperative game where we demonstrate how approximate inference can be successfully applied to instances that are too complex for exact computation. We discuss the relation of the KL control approach to other inference approaches to control.
1 Introduction
Stochastic optimal control is difficult in large or continuous state spaces, where standard Bellman-based computation and graphical value-function representations can become impractical. The paper recasts a class of these problems as KL minimization and graphical-model inference, enabling exact or approximate inference methods.
- Optimal control minimizes total future cost over actions and states, but computation is difficult because state spaces can be large and dynamics stochastic.
- Bellman equations compute optimal cost-to-go recursively, becoming a partial differential equation in continuous time.
- For high-dimensional or continuous systems, the state space is too large for direct Bellman computation, motivating function approximation and graphical-structure methods.
- Graphical structure may simplify computation, but it is generally not inherited by the value function, making graphical value-function representations inappropriate.
- The paper defines KL control problems by representing control as a trajectory distribution whose cost is a KL divergence, making optimization equivalent to inference in a dynamic Bayesian network.
- This formulation supports exact Junction Tree inference and approximate variational, belief propagation, cluster variation, generalized belief propagation, and MCMC methods.
- Unlike Todorov’s Bellman-equation identification, the paper identifies the control cost itself with a KL divergence, directly equating optimal control with graphical-model inference.
2 Control as KL minimization
The paper restricts stochastic optimal control to problems whose control cost is a KL divergence from free dynamics, turning optimal-control computation into graphical-model inference. The resulting formulation supports exact or approximate inference, includes path integral control as a special case, and exploits sparse structure in modular systems.
- KL control formulation: The controlled trajectory distribution is optimized against a free first-order Markov dynamics by minimizing their KL divergence plus state costs.The free dynamics may assign zero probability to physically impossible transitions, while the state cost remains arbitrary.
- KL control formulation: The KL minimization yields a normalized product of free dynamics and exponentiated costs, producing an optimal trajectory distribution that avoids high-cost states while staying close to q.Because q is first-order Markov, the resulting optimal distribution p is also first-order Markov.
- Graphical model inference: The optimal control at the current state is obtained from a marginal of the trajectory distribution, so its computation becomes a graphical-model inference problem.For chain-structured factors, the relevant marginal can be computed with backward messages.
- Graphical model inference: The optimal cost has a closed form as minus the log partition sum, equivalently an expectation of exponentiated state costs under the uncontrolled dynamics.This gives the optimal cost-to-go in terms of the known free dynamics q and state costs R.
- Path integral control: For continuous stochastic systems with quadratic control cost weighted by ν^-1, the KL formulation contains path integral control as a particular limit.The corresponding optimal cost-to-go is represented by an integral over trajectories.
- Graphical model inference: Modular systems make exact marginal computation intractable, but factorized uncontrolled dynamics and sparse component interactions create graphical structure exploitable by junction-tree or approximate inference methods.The paper identifies multi-agent systems and robot arms as examples where components are independently dynamic a priori but become coupled through task execution.
3 Stacking blocks (KL-blocks-world)
The KL-blocks-world models block stacking as stochastic control over graphical dynamics, then uses exact or approximate inference to compute action distributions and plans. Experiments show that CVM extends successful planning beyond the instances tractable by junction-tree inference or belief propagation.
- Problem formulation: The task models n block locations on a one-dimensional ring, with stack heights representing the state and one block moved to a neighboring location per iteration.Moves are represented by selecting a source location k_t and displacement l_t ∈ {-1, 0, 1}; forbidden moves remove blocks from empty stacks or exceed height m.
- Graphical model: The uncontrolled transition model is a mixture over source locations and displacements, with auxiliary variables enabling a compact graphical-model implementation.The graphical representation replaces a joint table with smaller interaction tables and supports efficient approximate inference.
- Control objectives: The paper studies both reaching a clamped goal state without state cost and minimizing entropy without a specified final state.The entropy objective favors concentrating all blocks on one stack when sufficient time is available.
- Exact inference: For a four-block tower shifted from position 1 to position 3, exact junction-tree inference produces symmetric action probabilities and a MAP plan that first unpacks, then rebuilds the tower.The MAP trajectory leaves one block at each of four locations at t = 4 before reconstructing the tower at the goal position.
- Approximate inference: Junction-tree inference becomes infeasible for larger instances, so the experiments compare belief propagation with the Cluster Variation Method using approximate inference.A four-block, six-location, horizon-11 instance requires approximately 15 Gbytes of memory under junction-tree inference.
- Approximate inference: BP solves all tested instances for n = 4 but fails to converge reliably for larger n, whereas CVM finds a correct plan in all runs at higher CPU cost and handles a large n = 8, m = 40, T = 80 instance.The large instance was too large for exact computation; CVM required approximately 1 hour per t iteration and 27 Mb of memory while yielding an apparently accurate solution.
4 Multi Agent cooperative game (KL-stag-hunt)
The KL-stag-hunt models multi-agent coordination as stochastic control on a grid, with rewards favoring individual hare capture or cooperative stag capture. Exact inference becomes infeasible as systems grow, while approximate methods—especially BP and CVM—recover useful strategies on larger instances with different convergence and computational trade-offs.
- Game formulation: The KL-stag-hunt is a finite-horizon grid game where agents move between locations and coordinate pairs to capture high-payoff stags.Agents may also capture hares individually for lower rewards; more than two agents do not increase a stag’s payoff.
- Game formulation: The state-dependent reward gives lower individual rewards for hares and higher rewards when at least two agents occupy a stag’s location.Hare rewards scale with the number of agents present, whereas stag rewards do not increase beyond two agents.
- Exact inference: For λ = 10, the optimal control is risk dominant and hunters pursue hares; for λ = 0.1, it is payoff dominant and hunters cooperate.The example uses a 5×5 grid, four corner hares, one central stag, N = 25, T = 4, Rs = −10 and Rh = −2.
- Graphical model: Exact joint-space inference becomes infeasible even for small agent counts because the joint state space scales as N^M.The factor-graph formulation is introduced to make approximate inference tractable.
- Approximate inference: For larger systems, BP produces efficient approximations with risk-dominant behavior at large λ and payoff-dominant behavior at sufficiently small λ.BP iterations increase as λ decreases, and BP may fail to converge for very small λ after 500 iterations.
5 Related work
The paper places KL control alongside inference-based control, reinforcement-learning policy updates, free-energy approaches, and EM methods. It distinguishes KL control by directly optimizing transition probabilities and obtaining optimal control from backward messages without policy-iteration steps.
- Inference-based control: Linear-quadratic control is equivalent to inference and can be solved as Kalman smoothing, while Kappen generalized the exponential relation to nonlinear continuous-time problems.The Bellman value function corresponds to the logarithm of a Kalman filter’s backward message.
- Inference-based control: Influence-diagram inference is efficient for single decisions but does not generalize efficiently to expected-reward multi-step tasks.The cited comparison attributes this limitation to the order of marginalization and optimization operations.
- Related approaches: KL control shares KL-divergence-based ideas with reinforcement-learning policy updating, including natural policy gradients and EM-type algorithms using relative entropy.These approaches are presented as related lines of research rather than as identical formulations.
- Related approaches: Path integral control is a special case of KL control for continuous-time systems with Gaussian noise.The paper presents this relationship as a limiting connection between the two formulations.
- Related approaches: Compared with TS, KL control uses a narrower reward form but is described as more efficient because backward messages directly provide optimal control without further iteration.TS permits arbitrary state-action rewards and alternates policy evaluation with policy improvement.
- Related approaches: Unlike TS, the KL approach’s complexity does not depend on whether the problem is time-dependent, and it optimizes transition probabilities directly rather than a policy.The paper contrasts this with TS’s fixed-policy message computation and subsequent policy optimization.
6 Discussion
The paper establishes an equivalence between a class of stochastic optimal control problems and graphical-model inference, enabling exact or approximate inference methods. Experiments show approximate controls for block stacking and multi-agent tasks, while the formulation has explicit scope limits.
- The KL control formulation makes stochastic optimal control equivalent to graphical-model inference, allowing exact or approximate inference methods to compute controls.
- The class excludes control problems lacking a q that satisfies the required KL relation and considers only model-based settings.
- Approximate inference successfully computes approximate controls in a block stacking task and a multi-agent cooperative task.
- In the KL-blocks-world, entropy minimization is harder than fixed-location stacking because the control must identify the optimal location.
- Standard BP is ineffective without a specified goal state, whereas generalized belief propagation using CVM produced apparently good optimal-control solutions.
- The formulation is particularly attractive for multi-agent problems because q factorizes over agents while rewards create interaction across agents.
- In the KL-stag-hunt game, BP approximates behavior in large systems where exact inference is infeasible.
- Strong penalties for deviations from the baseline policy make coordinated stag hunting sub-optimal, whereas unconstrained deviation makes it optimal.
A Boltzmann distribution
This section minimizes a KL-divergence objective over normalized distributions and derives a normalized Boltzmann-form solution. Substitution gives the minimum cost as the negative logarithm of the normalization constant.
- The objective is a KL divergence between a normalized distribution p(x) and a positive function ψ(x).
- Minimizing the objective with respect to p requires normalization, enforced with a Lagrange multiplier.
- Stationarity yields p(x) = ψ(x)/Z, where Z = exp(β + 1) normalizes the distribution.
- Substituting the minimizing distribution into the objective gives C = −log Z.
B Relation to continuous path integral model
The continuous-time construction connects controlled Gaussian dynamics to the KL formulation by taking the dt → 0 limit. The resulting divergence agrees with the earlier control equation.
- The controlled transition distribution is Gaussian with drift shifted by the control and covariance Ξ(x, t)dt.
- The derivation computes the transition mean and the quadratic control term involving g^T Ξ^-1g.
- The path-distribution entropy term appears in the trajectory-level KL expression.
- Taking the limit dt → 0 converts the KL divergence between controlled and uncontrolled path distributions into the continuous-time expression.
- The resulting continuous-time KL expression agrees with Equation (10).
C Cluster Variation Method
The Cluster Variation Method approximates the KL objective using overlapping cluster distributions and consistency constraints. Because its objective is non-convex, double-loop optimization provides guaranteed descent, while belief propagation is faster but not generally convergent.
- Cluster construction: The CVM replaces the full distribution with overlapping cluster distributions, each covering interactions among a small number of variables.
- Cluster construction: Figure 12 illustrates four clusters, five sub-clusters, and the non-convex CVM objective bounded by a convex function.
- Cluster construction: Clusters must contain every interaction term, while their intersections define sub-clusters used by the approximation.
- CVM objective: CVM minimizes an approximate KL divergence over cluster distributions subject to normalization and consistency constraints, using Möbius overcounting numbers.
- Optimization: Since overcounting numbers may be positive or negative, the CVM objective is non-convex.
- Optimization: The double-loop method constructs a convex upper bound and uses an inner convex optimization guaranteed to decrease the CVM objective to a local minimum.
- Optimization: Belief propagation solves equivalent fixed-point message-passing equations; it can be fast and accurate when convergent, but convergence is not guaranteed.