Policies as Probability Distributions

From Decision Rule to Probability Distribution

A policy answers one question: given state $s$, which action $a$ should the agent take? The simplest possible answer is a function.

DefinitionDeterministic Policy

A deterministic policy maps each state directly to one action:

$$a = \mu_\theta(s), \qquad \mu_\theta : \mathcal{S} \rightarrow \mathcal{A}$$

Given a fixed parameter vector $\theta$, the same state always produces the same action.

This is enough to define behavior, and it is often how a deterministic policy is actually implemented, for instance as a small multilayer perceptron mapping an observation vector directly to an action vector, with no sampling step and no probability distribution involved anywhere. But a purely deterministic map has a structural weakness already visible from an earlier note in this series: it cannot explore. If $\mu_\theta$ always returns the same action from the same state, the agent has no built-in mechanism for trying alternatives to find out whether they are better. Exploration has to be bolted on from outside, as explicit noise added after the fact.

DefinitionStochastic Policy

A stochastic policy defines a probability distribution over actions, conditioned on the state:

$$\pi_\theta(a \mid s)$$

Sampling an action from it is written $a \sim \pi_\theta(\cdot \mid s)$.

Exploration falls out of this definition for free: as long as $\pi_\theta$ assigns nonzero probability to more than one action, the agent will, by construction, sometimes try something other than what it currently believes is best. But, and this point is easy to blur, that is only one of two genuinely distinct reasons stochastic policies matter, and the other one turns out to be more consequential for what comes after this note.

RemarkTwo Different Reasons a Policy Should Be Stochastic

The first reason is behavioral: a distribution over actions naturally supports exploration, where a deterministic map does not. The second reason is mathematical, and has nothing to do with exploration at all: turning a decision into a distribution, rather than a single point, is what makes it possible to write down a smooth, differentiable objective over the policy’s own parameters, even when the underlying choice is discrete, or when the quantity being optimized (return) is not differentiable in the action itself. This second reason is why policy-gradient methods need $\pi_\theta(a \mid s)$ specifically, as a genuine probability distribution with a density or mass function, and not merely “some randomness added to a deterministic policy.” It is the subject of the entire log-derivative trick worked out in the Policy Gradient note.

The concrete shape $\pi_\theta(a \mid s)$ takes depends entirely on the action space it has to cover, and, as with the value-based methods contrasted earlier in this series, discrete and continuous action spaces call for genuinely different constructions.

Discrete Actions: The Categorical Policy

DefinitionCategorical Policy

For a discrete action space $\mathcal{A} = \{a_1, \dots, a_n\}$, a categorical policy outputs a probability vector

$$\pi_\theta(\cdot \mid s) = \begin{bmatrix} p_1 \\ \vdots \\ p_n \end{bmatrix}, \qquad p_i = \pi_\theta(a_i \mid s), \qquad \sum_{i=1}^n p_i = 1.$$

In practice this is implemented exactly like a classifier: a network outputs $n$ real-valued logits, and a softmax converts them into the probability vector above.

If the sampled action is $a_i$, its likelihood is $\pi_\theta(a_i \mid s) = p_i$, and, the quantity that will matter for everything downstream, its log likelihood is $\log \pi_\theta(a_i \mid s) = \log p_i$.

Continuous Actions: The Diagonal Gaussian Policy

A categorical distribution has no meaningful analogue once the action space becomes a continuum. There is no way to output “a probability for every real number.” The standard construction instead outputs the parameters of a continuous distribution.

DefinitionDiagonal Gaussian Policy

For a $k$-dimensional continuous action $a \in \mathbb{R}^k$, a diagonal Gaussian policy is

$$\pi_\theta(a \mid s) = \mathcal{N}\big(\mu_\theta(s),\ \operatorname{diag}(\sigma^2)\big)$$

with mean vector $\mu_\theta(s)$ and covariance matrix restricted to be diagonal, $\Sigma = \operatorname{diag}(\sigma_1^2, \dots, \sigma_k^2)$.

RemarkWhy Diagonal

Restricting $\Sigma$ to be diagonal is a genuine modeling choice, not a mathematical necessity. It says the action dimensions are treated as conditionally independent given the state, $a_i \sim \mathcal{N}(\mu_i, \sigma_i^2)$ for each $i$ separately. A full covariance matrix could in principle capture correlations between action dimensions (steering angle and speed moving together, say), but at the cost of $O(k^2)$ parameters and a matrix that must be kept positive-definite throughout training. The diagonal restriction trades that expressiveness for $O(k)$ parameters and a positivity constraint that only ever has to be enforced coordinate-by-coordinate.

RemarkAn Unconstrained Parameterization for a Constrained Quantity

A standard deviation must satisfy $\sigma_i > 0$, but a neural network’s raw output is an unconstrained real number. Rather than clip or otherwise hard-constrain the network’s output, the standard move is to have the network output $\log \sigma_i \in \mathbb{R}$ instead, and recover $\sigma_i = \exp(\log \sigma_i)$. Since $\exp: \mathbb{R} \to \mathbb{R}_{>0}$ is a bijection, this reparameterization enforces positivity exactly, for free, without ever touching the optimizer’s unconstrained parameter space.

Sampling from this policy uses the same trick in reverse: rather than sampling directly from a distribution whose parameters depend on $\theta$, first sample parameter-free noise, then apply a deterministic, $\theta$-dependent transformation to it.

DefinitionSampling via Reparameterization

Sample $z \sim \mathcal{N}(0, I)$, then set

$$a = \mu_\theta(s) + \sigma \odot z, \qquad\text{i.e. coordinate-wise } a_i = \mu_i(s) + \sigma_i z_i,$$

where $\odot$ denotes element-wise multiplication.

PropositionThe Reparameterization Is Correct

The $a$ constructed above satisfies $a \sim \mathcal{N}(\mu_\theta(s), \operatorname{diag}(\sigma^2))$.

Proof

Each coordinate $a_i = \mu_i(s) + \sigma_i z_i$ is an affine transformation of the standard normal variable $z_i \sim \mathcal{N}(0,1)$. An affine transformation of a Gaussian is Gaussian, with mean and variance transforming as $\mathbb{E}[a_i] = \mu_i(s) + \sigma_i \mathbb{E}[z_i] = \mu_i(s)$ and $\operatorname{Var}(a_i) = \sigma_i^2 \operatorname{Var}(z_i) = \sigma_i^2$, so $a_i \sim \mathcal{N}(\mu_i(s), \sigma_i^2)$. The coordinates of $z$ are independent by construction, and an affine, coordinate-wise transformation of independent variables preserves independence, so the $a_i$ are independent as well, which is exactly $a \sim \mathcal{N}(\mu_\theta(s), \operatorname{diag}(\sigma^2))$.

The log likelihood follows from writing out the diagonal Gaussian density and taking its logarithm.

DefinitionLog Likelihood of a Diagonal Gaussian Policy

$$\log \pi_\theta(a \mid s) = -\frac{1}{2} \sum_{i=1}^{k} \left[ \left(\frac{a_i - \mu_i}{\sigma_i}\right)^2 + 2\log\sigma_i + \log(2\pi) \right]$$

Actions close to the mean have higher log likelihood; actions far from the mean have lower log likelihood; and a larger $\sigma_i$ both widens the distribution and, visible directly in the formula, flattens the log-likelihood’s sensitivity to how far $a_i$ lands from the mean, in exchange for a constant $-\log\sigma_i$ penalty. Larger $\sigma_i$ trades precision for exploration in dimension $i$, and the log-likelihood formula prices that trade-off explicitly.

Log Likelihood: The Object Policies Actually Optimize

Both constructions above end in the same place: not a probability, but a log probability. That is not a stylistic preference.

RemarkWhy Log, Specifically

Writing $\log \pi_\theta(a \mid s)$ rather than $\pi_\theta(a \mid s)$ does two things at once. First, it turns products of probabilities into sums: a trajectory’s probability under a policy factorizes as a product of per-step probabilities, so its log probability is a sum of per-step log probabilities, which is both more numerically stable for long sequences of small numbers and dramatically easier to differentiate term by term. Second, and more fundamentally, the log likelihood is the exact object the log-derivative trick operates on: $\nabla_\theta \log \pi_\theta(a \mid s)$ is precisely the score function that turns “the return of a sampled action” into a usable gradient signal, without ever needing to differentiate through the environment. That identity, why the gradient of a log probability is the right thing to multiply a return by, is substantial enough to deserve its own derivation, and gets one in full in the Policy Gradient note.

Comparison

Policy type	Action space	Distribution	Main output
Deterministic policy	continuous or discrete	none	direct action $a = \mu_\theta(s)$
Categorical policy	discrete	categorical	action probabilities $\pi_\theta(a_i \mid s)$
Diagonal Gaussian policy	continuous	Gaussian	mean $\mu_\theta(s)$ and standard deviation $\sigma$

References

• OpenAI Spinning Up: Part 1, Key Concepts in RL