Control as inference is a probabilistic framework that recasts ♟️ Reinforcement Learning through the lens of 🧬 Variational Inference in a 🪩 Probabilistic Graphical Model. The key idea of this framework is that unlike other optimal methods, this can model suboptimality and give us "soft" policies that explain stochastic behavior.

The core idea is that a trajectory can be augmented with binary optimality variables , where we define

This optimality variable allows us to describe the probability of optimal behavior: that is, the probability of a trajectory given optimality is

In other words, it's the probability that the trajectory can actually happen, , weighted by its exponentiated reward.

Planning via Inference#

Within this framework, we can compute multiple useful quantities.

1. Backward messages , ie how likely we are to be optimal at and onwards.
2. Optimal policy .
3. Forward messages , ie probability of reaching a state , given optimal behavior.

Note that this terminology takes inspiration from the ☂️ Hidden Markov Model > Forward-Backward Algorithm, with a similar message passing technique.

Backward Message#

The backward message is . Expanding this out gives

and the first term, the backward state message, is

We can assume that (actions taking without optimality in mind) is uniform, and noting that the first term is the previous backward message, we have

The backward message passing thus amounts to alternating between computing the two messages starting from down to : 1.

Though the above equations don't seem very intuitive, we can recast them in the form of 💎 Value Iteration if we let

Then, the above steps are the following: 1.

Note that rather than computing the maximum, our log of the exponentiation finds a "soft" maximum (not to be confused with softmax).

Policy Derivation#

Using the backward messages, we can derive, through 🪙 Bayes' Theorem,

Note that in context of our value functions, this is equivalent to

which essentially gives better actions an exponentially higher probability.

Forward Message#

A forward message gives us more insight into the trajectories we'll reach given optimality. Some lengthy algebra gives us

where

More importantly, if we now compute the state marginal under optimality, we find

Intuitively, this tells us that the state distribution is the intersection of states with high probability of reaching the goal (backward) and states with high probability of originating from the initial state (forward).

Fixing Optimistic Transitions#

Unfortunately, an important problem with the above approach is the equation

which describes an "optimistic" transition (as opposed to a normal expectation, in 💎 Value Iteration). In taking this "soft" max over the future states, it is based on the assumption that the future state is also optimal, which is up to the stochastic environment transitions and not within our control. That is, it is confusing two questions into one:

1. Given that we obtained high reward, what was our action probability? This is the question we want to answer.
2. Given that we obtained high reward, what was our transition probability? This is not controllable, and thus should not factor into our planning.

To address this, we'll harness 🧬 Variational Inference and find another distribution that's close to but with the actual dynamics (unaffected by optimality). This is an approximation problem: with and , we see that our problem reduces to finding to approximate .

First, to enforce the dynamics, we define

where , our policy, is the only part of the distribution we can control. Applying the 💼 Evidence Lower Bound, we have

Tightening the bound by maximizing this quantity is thereby equivalent to maximizing reward and action entropy. From here, it can be shown that the policy that maximizes this value is

with value functions

These definitions for value functions and policy can be immediately substituted into 🚀 Q-Learning and 🚓 Policy Gradients for their soft optimality variants, bringing the benefits of improved exploration, easier fine-tuning, and stronger robustness. Using a similar idea gives us 🎲 Entropy Regularization, which forms the basis of 🪶 Soft Actor-Critic.