Ahmed Haroon

From Control to Inference

Ordinary optimal control chooses actions that maximize cumulative reward:

a_{1:T}^* =\arg\max_{a_{1:T}} \sum_{t=1}^{T}r(s_t,a_t),

subject to the environment dynamics. This produces a hard distinction between optimal and non-optimal behavior.

But several behaviors may all be good, demonstrations may contain mistakes, and a controller may benefit from remaining stochastic. Control as inference replaces hard optimality with a probabilistic event.

Introduce a binary optimality variable $O_t$ at every time step and define

\boxed{ p(O_t=1\mid s_t,a_t) \propto \exp\left(\frac{r(s_t,a_t)}{\alpha}\right) },

where $\alpha>0$ is a temperature. High-reward state-action pairs make optimality more likely, without forcing lower-reward actions to have zero probability.

Before conditioning, let the trajectory prior be

p(\tau) =p(s_1) \prod_{t=1}^{T}p_0(a_t\mid s_t) \prod_{t=1}^{T-1}p(s_{t+1}\mid s_t,a_t),

where $p_0(a\mid s)$ is an action prior. The joint distribution is

p(\tau,O_{1:T}=1) =p(\tau) \prod_{t=1}^{T} \exp\left(\frac{r(s_t,a_t)}{\alpha}\right).

Therefore,

\boxed{ p(\tau\mid O_{1:T}=1) \propto p(\tau) \exp\left( \frac1\alpha\sum_{t=1}^{T}r(s_t,a_t) \right) }.

The prior keeps trajectories physically plausible. The exponential-return factor assigns greater posterior probability to higher-return trajectories.

For equally likely trajectories $\tau_1$ and $\tau_2$ ,

\frac{p(\tau_1\mid O=1)}{p(\tau_2\mid O=1)} =\exp\left( \frac{R(\tau_1)-R(\tau_2)}{\alpha} \right).

A return advantage of $\Delta R$ multiplies posterior odds by $e^{\Delta R/\alpha}$ .

Backward Inference

Define the state-action backward message

\beta_t(s_t,a_t) =p(O_{t:T}=1\mid s_t,a_t).

It asks how likely future behavior is to be optimal after taking $a_t$ in $s_t$ . Also define

\beta_t(s_t)=p(O_{t:T}=1\mid s_t).

The messages satisfy

\boxed{ \beta_t(s,a) =\exp\left(\frac{r(s,a)}{\alpha}\right) \mathbb E_{s'\sim p(\cdot\mid s,a)} [\beta_{t+1}(s')] },

and

\boxed{ \beta_t(s) =\mathbb E_{a\sim p_0(\cdot\mid s)} [\beta_t(s,a)] }.

Take scaled logarithms so the messages have reward units:

Q_t(s,a)=\alpha\log\beta_t(s,a), \qquad V_t(s)=\alpha\log\beta_t(s).

This produces the exact inference recursions

\boxed{ Q_t(s,a) =r(s,a) +\alpha\log \mathbb E_{s'\sim p(\cdot\mid s,a)} \left[ \exp\left(\frac{V_{t+1}(s')}{\alpha}\right) \right] },

\boxed{ V_t(s) =\alpha\log \int p_0(a\mid s) \exp\left(\frac{Q_t(s,a)}{\alpha}\right)da }.

See derivation

Starting from the state-action message,

\beta_t(s,a) =e^{r(s,a)/\alpha} \mathbb E_{s'}[\beta_{t+1}(s')],

take logs and multiply by $\alpha$ :

\begin{aligned} Q_t(s,a) &=\alpha\log\beta_t(s,a)\\ &=r(s,a) +\alpha\log\mathbb E_{s'} [\beta_{t+1}(s')]\\ &=r(s,a) +\alpha\log\mathbb E_{s'} \left[e^{V_{t+1}(s')/\alpha}\right]. \end{aligned}

For the state message,

\begin{aligned} V_t(s) &=\alpha\log \mathbb E_{a\sim p_0}[\beta_t(s,a)]\\ &=\alpha\log \int p_0(a\mid s)e^{Q_t(s,a)/\alpha}da. \end{aligned}

For discrete actions with a uniform prior absorbed into the base measure,

V_t(s) =\alpha\log\sum_a \exp\left(\frac{Q_t(s,a)}{\alpha}\right).

This log-sum-exp is a smooth maximum. If $M=\max_iQ_i$ , then

\boxed{ M \leq \alpha\log\sum_i e^{Q_i/\alpha} \leq M+\alpha\log n }.

Therefore,

\lim_{\alpha\to0^+} \alpha\log\sum_i e^{Q_i/\alpha} =\max_iQ_i.

For deterministic dynamics $s'=f(s,a)$ , the Q-backup becomes

Q_t(s,a)=r(s,a)+V_{t+1}(f(s,a)),

so backward inference is soft dynamic programming.

The Posterior Policy

The inferred policy is the posterior action distribution conditioned on future optimality:

\pi^*(a_t\mid s_t) =p(a_t\mid s_t,O_{t:T}=1).

Bayes' rule gives

\pi^*(a\mid s) =p_0(a\mid s) \frac{\beta_t(s,a)}{\beta_t(s)}.

Using the log-messages,

\boxed{ \pi^*(a\mid s) =p_0(a\mid s) \exp\left( \frac{Q_t(s,a)-V_t(s)}{\alpha} \right) }.

If $A_t(s,a)=Q_t(s,a)-V_t(s)$ , then

\pi^*(a\mid s) =p_0(a\mid s)e^{A_t(s,a)/\alpha}.

The policy is normalized because

e^{V_t(s)/\alpha} =\int p_0(a\mid s)e^{Q_t(s,a)/\alpha}da.

For two actions,

\frac{\pi^*(a_1\mid s)}{\pi^*(a_2\mid s)} =\frac{p_0(a_1\mid s)}{p_0(a_2\mid s)} \exp\left( \frac{Q(s,a_1)-Q(s,a_2)}{\alpha} \right).

Small $\alpha$ concentrates probability on maximizing actions. Large $\alpha$ makes the policy approach its prior:

\lim_{\alpha\to0^+}\pi^*(a\mid s) =\text{greedy policy},

\lim_{\alpha\to\infty}\pi^*(a\mid s) =p_0(a\mid s).

The action prior has three equivalent interpretations:

a default behavior before conditioning on optimality;
an implicit reward $\alpha\log p_0(a\mid s)$ ;
a reference distribution that the learned policy is penalized for leaving.

Forward Messages and Occupancy

Backward messages say whether a state can lead to high future reward. Forward messages say whether it is reachable from the initial distribution while remaining optimal so far.

Define the unnormalized forward message

\widetilde f_t(s_t) =p(s_t,O_{1:t-1}=1).

It follows the recursion

\boxed{ \widetilde f_{t+1}(s') =\int \widetilde f_t(s) p_0(a\mid s) e^{r(s,a)/\alpha} p(s'\mid s,a) \,da\,ds }.

The state posterior conditioned on optimality over the entire trajectory is

\boxed{ p(s_t\mid O_{1:T}=1) \propto \widetilde f_t(s_t)\beta_t(s_t) }.

This is a forward-backward intersection:

\underbrace{\widetilde f_t(s_t)}_{\text{reachable from the start}} \times \underbrace{\beta_t(s_t)}_{\text{can lead to high reward}}.

A state may be promising but unreachable, or reachable but unable to lead to a good outcome. The conditioned occupancy emphasizes states satisfying both conditions.

The Optimism Problem

Exact inference conditions every random variable on optimality, including stochastic next states. The posterior transition becomes

\boxed{ p(s'\mid s,a,O_{t:T}=1) =\frac{ p(s'\mid s,a)\beta_{t+1}(s') }{ \mathbb E_{\bar s\sim p(\cdot\mid s,a)} [\beta_{t+1}(\bar s)] } }.

Favorable next states receive greater posterior probability. Unless $\beta_{t+1}$ is constant,

p(s'\mid s,a,O_{t:T}=1) \neq p(s'\mid s,a).

The same optimism appears in the exact-inference backup:

Q_{\mathrm{exact}}(s,a) =r(s,a)+\alpha\log \mathbb E_{s'}[e^{V(s')/\alpha}].

By Jensen's inequality,

\boxed{ \alpha\log\mathbb E[e^{V/\alpha}] \geq\mathbb E[V] }.

For small $\alpha$ , even a rare favorable outcome can dominate the log-sum-exp.

Suppose an action leads to values $10$ and $0$ with equal probability and $\alpha=1$ . Ordinary expected value is

\mathbb E[V]=5,

while exact inference gives

\log\left(\frac12e^{10}+\frac12e^0\right) \approx9.31.

This posterior statement is valid:

\text{given a high-reward trajectory, a favorable transition probably occurred}.

But this control assumption is invalid:

\text{choosing the action makes the favorable transition more likely}.

A controller may change its action distribution. It cannot change the environment dynamics. We therefore need an inference approximation that keeps transitions fixed.

Control as Variational Inference

Choose a variational trajectory distribution parameterized only by a policy:

\boxed{ q_\pi(\tau) =p(s_1) \prod_{t=1}^{T}\pi(a_t\mid s_t) \prod_{t=1}^{T-1}p(s_{t+1}\mid s_t,a_t) }.

It retains the real initial-state and transition distributions. Only the action prior changes:

p_0(a_t\mid s_t) \longrightarrow \pi(a_t\mid s_t).

Maximizing the ELBO gives

\boxed{ \alpha\mathcal L(q_\pi) =\mathbb E_{\tau\sim q_\pi} \left[ \sum_{t=1}^{T} \left( r(s_t,a_t) -\alpha\log \frac{\pi(a_t\mid s_t)}{p_0(a_t\mid s_t)} \right) \right] }.

See derivation

The variational lower bound is

\mathcal L(q_\pi) =\mathbb E_{q_\pi} \left[ \log p(\tau,O_{1:T}=1)-\log q_\pi(\tau) \right].

Expand the joint model:

\begin{aligned} \log p(\tau,O=1) &=\log p(s_1) +\sum_{t=1}^{T-1}\log p(s_{t+1}\mid s_t,a_t)\\ &\quad+\sum_{t=1}^{T}\log p_0(a_t\mid s_t) +\frac1\alpha\sum_{t=1}^{T}r(s_t,a_t). \end{aligned}

The variational log-density is

\log q_\pi(\tau) =\log p(s_1) +\sum_{t=1}^{T-1}\log p(s_{t+1}\mid s_t,a_t) +\sum_{t=1}^{T}\log\pi(a_t\mid s_t).

Thus

\mathcal L(q_\pi) =\mathbb E_{q_\pi} \left[ \sum_t \left( \frac{r(s_t,a_t)}\alpha +\log p_0(a_t\mid s_t) -\log\pi(a_t\mid s_t) \right) \right].

Multiplying by $\alpha$ yields the stated objective.

At a fixed state,

\mathbb E_{a\sim\pi} \left[ \log\frac{\pi(a\mid s)}{p_0(a\mid s)} \right] =D_{\mathrm{KL}} \left( \pi(\cdot\mid s)\|p_0(\cdot\mid s) \right).

Therefore variational control maximizes reward while penalizing deviation from a default behavior:

\boxed{ J(\pi) =\mathbb E_{\tau\sim\pi} \left[ \sum_t \left( r(s_t,a_t) -\alpha D_{\mathrm{KL}} (\pi(\cdot\mid s_t)\|p_0(\cdot\mid s_t)) \right) \right] }.

Maximum-Entropy Reinforcement Learning

For a uniform action prior,

-D_{\mathrm{KL}}(\pi\|p_0) =\mathcal H(\pi)-\text{constant}.

The variational objective becomes

\boxed{ J(\pi) =\mathbb E_{\tau\sim\pi} \left[ \sum_{t=0}^{\infty}\gamma^t \left( r(s_t,a_t) +\alpha\mathcal H(\pi(\cdot\mid s_t)) \right) \right] }.

Equivalently, using sampled actions,

J(\pi) =\mathbb E_{\tau\sim\pi} \left[ \sum_{t=0}^{\infty}\gamma^t \left( r(s_t,a_t)-\alpha\log\pi(a_t\mid s_t) \right) \right].

Entropy is not rewarded independently of return. Low-value actions remain exponentially unlikely. The objective mainly spreads probability among competitive actions, avoids premature collapse, and preserves multiple good solutions.

Soft policy evaluation

For a fixed policy,

\boxed{ Q^\pi(s,a) =r(s,a)+\gamma \mathbb E_{s'\sim p(\cdot\mid s,a)}[V^\pi(s')] },

\boxed{ V^\pi(s) =\mathbb E_{a\sim\pi(\cdot\mid s)} \left[ Q^\pi(s,a) -\alpha\log \frac{\pi(a\mid s)}{p_0(a\mid s)} \right] }.

The transition expectation is ordinary. The action value includes reward and future soft value; the state value also includes the policy's entropy or KL cost.

Soft policy improvement

Given $Q^{\pi_{\mathrm{old}}}$ , improve the policy by solving

\max_{\pi'} \left{ \mathbb E_{a\sim\pi'}[Q^{\pi_{\mathrm{old}}}(s,a)] -\alpha D_{\mathrm{KL}}(\pi'\|p_0) \right}.

The solution is

\boxed{ \pi_{\mathrm{new}}(a\mid s) =\frac{ p_0(a\mid s) \exp(Q^{\pi_{\mathrm{old}}}(s,a)/\alpha) }{ \int p_0(\bar a\mid s) \exp(Q^{\pi_{\mathrm{old}}}(s,\bar a)/\alpha)d\bar a } }.

See derivation

Define

Z_Q(s)=\int p_0(a\mid s)e^{Q(s,a)/\alpha}da

and

\pi_Q(a\mid s) =\frac{p_0(a\mid s)e^{Q(s,a)/\alpha}}{Z_Q(s)}.

Then

Q(s,a)+\alpha\log p_0(a\mid s) =\alpha\log\pi_Q(a\mid s)+\alpha\log Z_Q(s).

Substitute into the improvement objective:

\begin{aligned} &\mathbb E_{\pi'}[Q] -\alpha D_{\mathrm{KL}}(\pi'\|p_0)\\ &=\alpha\log Z_Q(s) -\alpha D_{\mathrm{KL}}(\pi'\|\pi_Q). \end{aligned}

KL divergence is minimized at zero when

Therefore $\pi_Q$ is the improving policy.

Soft Bellman optimality

At the optimal policy,

\boxed{ Q^*(s,a) =r(s,a)+\gamma\mathbb E_{s'}[V^*(s')] },

\boxed{ V^*(s) =\alpha\log \int p_0(a\mid s)e^{Q^*(s,a)/\alpha}da },

and

\boxed{ \pi^*(a\mid s) =p_0(a\mid s) \exp\left( \frac{Q^*(s,a)-V^*(s)}{\alpha} \right) }.

The difference from exact inference is crucial:

Q_{\mathrm{exact}} =r+\alpha\log\mathbb E_{s'}[e^{V(s')/\alpha}],

Q_{\mathrm{variational}} =r+\gamma\mathbb E_{s'}[V(s')].

Exact inference softens and reweights stochastic transitions. Variational maximum-entropy control keeps transitions fixed and softens only action selection.

\textbf{Soft Policy Iteration}

\begin{aligned} &\text{Initialize a policy }\pi_0.\\[0.25em] &\text{Repeat:}\\[0.25em] &\qquad \text{evaluate }Q^{\pi_k}\text{ using}\\[0.25em] &\qquad\quad Q^{\pi_k}(s,a) \leftarrow r(s,a)+\gamma\mathbb E_{s'} \mathbb E_{a'\sim\pi_k} \left[ Q^{\pi_k}(s',a') -\alpha\log\frac{\pi_k(a'\mid s')}{p_0(a'\mid s')} \right],\\[0.25em] &\qquad \text{improve with}\\[0.25em] &\qquad\quad \pi_{k+1}(a\mid s) \propto p_0(a\mid s)e^{Q^{\pi_k}(s,a)/\alpha}. \end{aligned}

For small discrete problems, these steps can be exact. For continuous actions, the log-partition integral and exact soft policy are usually intractable. This motivates a learned actor.

Soft Actor-Critic

SAC approximates soft policy iteration with:

a stochastic actor $\pi_\phi$ ;
twin critics $Q_{\theta_1}$ and $Q_{\theta_2}$ ;
slowly updated target critics;
experience replay;
reparameterized action samples.

Soft critic target

For a replay transition $(s,a,r,s',d)$ , sample

a'\sim\pi_\phi(\cdot\mid s').

The target is

\boxed{ y =r+\gamma(1-d) \left[ \min_{i\in\{1,2\}}Q_{\bar\theta_i}(s',a') -\alpha\log\pi_\phi(a'\mid s') \right] }.

The minimum reduces positive critic bias. The negative log-probability is the sampled entropy bonus.

Each critic minimizes

\boxed{ \mathcal L_Q(\theta_i) =\mathbb E_\mathcal D \left[ \left(Q_{\theta_i}(s,a)-\operatorname{stopgrad}[y]\right)^2 \right] }.

The targets move slowly through Polyak averaging:

\bar\theta_i \leftarrow \rho\bar\theta_i+(1-\rho)\theta_i, \qquad \rho\approx1.

Squashed reparameterized actions

The actor predicts a mean and standard deviation. Sample

\epsilon\sim\mathcal N(0,I),

u=\mu_\phi(s)+\sigma_\phi(s)\odot\epsilon,

and squash to the action bounds:

\boxed{a=\tanh(u)}.

This turns stochastic sampling into a differentiable function of parameter-free noise. Gradients can flow from the critic through the action into the actor.

The transformation changes the action density. Since

\frac{\partial a_j}{\partial u_j} =1-\tanh^2(u_j) =1-a_j^2,

the change-of-variables formula gives

\boxed{ \log\pi_\phi(a\mid s) =\log\mathcal N (u;\mu_\phi(s),\sigma_\phi^2(s)) -\sum_j\log(1-a_j^2) }.

In practice, we add a small $\varepsilon$ or use an algebraically stable expression. Without this correction, the entropy term describes the unsquashed Gaussian rather than the bounded action distribution actually sent to the environment.

Actor loss

Soft policy improvement fits the actor to the exponentiated-Q distribution. The resulting loss is

\boxed{ \mathcal L_\pi(\phi) =\mathbb E_{ s\sim\mathcal D,\, a\sim\pi_\phi } \left[ \alpha\log\pi_\phi(a\mid s) -\min_{i\in\{1,2\}}Q_{\theta_i}(s,a) \right] }.

See derivation

Let the target soft policy be

\pi_Q(a\mid s) =\frac{e^{Q(s,a)/\alpha}}{Z(s)}.

Fit the actor by minimizing reverse KL:

\begin{aligned} D_{\mathrm{KL}}(\pi_\phi\|\pi_Q) &=\mathbb E_{a\sim\pi_\phi} \left[ \log\pi_\phi(a\mid s) -\frac{Q(s,a)}\alpha +\log Z(s) \right]. \end{aligned}

Multiply by $\alpha$ :

\alpha D_{\mathrm{KL}}(\pi_\phi\|\pi_Q) =\mathbb E_{\pi_\phi} [\alpha\log\pi_\phi(a\mid s)-Q(s,a)] +\alpha\log Z(s).

Dropping it yields the SAC actor loss.

The two terms exert opposing pressures:

\underbrace{-Q(s,a)}_{\text{choose high-value actions}} + \underbrace{\alpha\log\pi(a\mid s)}_{\text{avoid excessive concentration}}.

Automatic temperature tuning

A fixed $\alpha$ is sensitive to reward scale and action dimension. SAC can learn it to match a target entropy $\mathcal H_{\mathrm{tgt}}$ :

\boxed{ \mathcal L_\alpha =\mathbb E \left[ -\alpha \left( \log\pi_\phi(a\mid s) +\mathcal H_{\mathrm{tgt}} \right) \right] }.

Parameterize

\alpha=e^\eta, \qquad \eta=\log\alpha,

to guarantee positivity. If entropy is below the target, optimization increases $\alpha$ ; if it is above the target, $\alpha$ decreases. The actor's log-probability is detached during the temperature update.

\textbf{Soft Actor-Critic}

\begin{aligned} &\text{Initialize }\pi_\phi,Q_{\theta_1},Q_{\theta_2}, \text{ target critics, and replay }\mathcal D.\\[0.25em] &\text{For each environment step:}\\[0.25em] &\qquad a_t\sim\pi_\phi(\cdot\mid s_t),\\[0.25em] &\qquad \text{execute }a_t\text{ and store the transition},\\[0.25em] &\qquad (s,a,r,s',d)\sim\mathcal D,\\[0.25em] &\qquad a'\sim\pi_\phi(\cdot\mid s'),\\[0.25em] &\qquad y\leftarrow r+\gamma(1-d) \left[ \min_iQ_{\bar\theta_i}(s',a') -\alpha\log\pi_\phi(a'\mid s') \right],\\[0.25em] &\qquad \text{fit both critics to the detached target }y,\\[0.25em] &\qquad a\sim\pi_\phi(\cdot\mid s)\text{ by reparameterization},\\[0.25em] &\qquad \phi\leftarrow\arg\min_\phi \mathbb E[\alpha\log\pi_\phi(a\mid s)-\min_iQ_{\theta_i}(s,a)],\\[0.25em] &\qquad \text{optionally update }\alpha,\\[0.25em] &\qquad \text{Polyak-update the target critics}. \end{aligned}

Boundaries and Special Cases

Boltzmann exploration is not maximum-entropy RL. An ordinary Q-learning agent may select actions with

\pi_{\mathrm{beh}}(a\mid s) \propto e^{Q_{\mathrm{ordinary}}(s,a)/\alpha}

while still training the critic with

y=r+\gamma\max_{a'}Q(s',a').

Here softmax is only an exploration rule. In maximum-entropy RL, entropy changes the objective, value function, critic target, and policy-improvement operator:

\boxed{ \text{softmax action selection} \neq \text{maximum-entropy reinforcement learning} }.

Non-uniform priors. The general KL-regularized objective is

J(\pi) =\mathbb E_\pi \left[ \sum_t\gamma^t \left( r(s_t,a_t) -\alpha D_{\mathrm{KL}} (\pi(\cdot\mid s_t)\|p_0(\cdot\mid s_t)) \right) \right].

The prior can represent a previous policy, safe controller, demonstration policy, skill, or reference model. The learned policy deviates only when reward justifies the KL cost.

Ordinary reinforcement learning. As temperature approaches zero,

V^*(s) =\alpha\log\int p_0(a\mid s)e^{Q^*(s,a)/\alpha}da \longrightarrow \max_{a:p_0(a\mid s)>0}Q^*(s,a).

Therefore,

\boxed{ \lim_{\alpha\to0^+} \text{maximum-entropy RL} =\text{ordinary optimal control} }.

At the other extreme,

\lim_{\alpha\to\infty}\pi^*(a\mid s) =p_0(a\mid s).

Temperature interpolates between greedy reward maximization and following the prior.

Practical Diagnosis

Symptom	Likely cause	Useful response
Critic values are too optimistic	Transition inference or critic overestimation	Keep real dynamics fixed and use twin critics
SAC entropy collapses early	Temperature is too small	Tune or learn $\alpha$ and inspect target entropy
SAC remains random despite learning	Entropy pressure dominates reward	Check reward scale and temperature
Actions reach bounds but log-probabilities look normal	Missing tanh Jacobian correction	Compute density after the action transformation
Q-values explode	Target graph, terminal mask, or reward scale is wrong	Detach targets and verify bootstrapping
Temperature moves in the wrong direction	Sign convention is inconsistent	Check entropy against $\mathcal H_{\mathrm{tgt}}$
Actor improves Q but not environment return	Critic is being exploited	Monitor realized return and use conservative critics

The full derivation can be remembered as

p(O_t=1\mid s_t,a_t) \propto e^{r(s_t,a_t)/\alpha}

\Downarrow

p(\tau\mid O_{1:T}=1) \propto p(\tau)e^{R(\tau)/\alpha}

\Downarrow\quad\text{restrict inference to real dynamics}

\mathbb E_\pi \left[ \sum_t r(s_t,a_t) -\alpha\log\frac{\pi(a_t\mid s_t)}{p_0(a_t\mid s_t)} \right]

\Downarrow\quad p_0\text{ uniform}

\mathbb E_\pi \left[ \sum_t r(s_t,a_t)+\alpha\mathcal H(\pi(\cdot\mid s_t)) \right].

Control as inference turns rewards into evidence about optimal behavior. Variational control prevents that inference from cheating through favorable dynamics. Soft Actor-Critic then approximates the resulting soft policy iteration with replay, twin critics, and a reparameterized stochastic actor.