Section 15.1: Direct policy optimization; stochastic policies

A Careful Control Loop

Value-based methods ask, "Which action has the highest estimated value?" Policy-gradient methods ask a different question: "How should the parameters of the action distribution move so future sampled behavior becomes more successful?" That shift is important in embodied control because the action may be continuous, multimodal, or constrained by hardware limits.

The object we optimize is the expected discounted return:

$$J(\theta)=\mathbb{E}_{\tau\sim\pi_\theta}\left[\sum_{t=0}^{T-1}\gamma^t r_t\right].$$

Here $\theta$ are policy parameters, $\tau$ is a trajectory of observations, actions, and rewards, $\pi_\theta(a_t \mid o_t)$ is the probability or density assigned to an action, and $\gamma$ discounts delayed consequences. Direct optimization means we adjust $\theta$ to increase $J(\theta)$ rather than first learning a separate action-value function and acting greedily from it.

Theory

The policy is a distribution, not a lookup table. For a discrete mobile robot action set, $\pi_\theta$ might be a softmax over forward, left, right, and stop. For a manipulation controller, it might be a Gaussian over joint velocity commands with a learned mean and standard deviation. In both cases, the policy must expose two things for training: the sampled action and the log probability of that exact action.

The central intuition is credit assignment through sampling. If a robot nudges a drawer handle and the drawer opens later, the update cannot differentiate through the drawer physics, contact impulses, camera exposure, or environment reset. It can still differentiate through the policy's own probability of the sampled motion. That is the foothold used by REINFORCE and PPO.

Worked Example

Code Fragment 1 below shows the smallest useful contrast: a deterministic controller always chooses the largest probability action, while a stochastic controller samples from the whole distribution and records the probability of what it did.

# Compare greedy action selection with stochastic sampling.
# The sampled action keeps a log probability, which policy gradients need.
import math
import random

random.seed(7)
actions = ["move_left", "move_right", "stop"]
logits = [0.2, 1.4, -0.7]

exp_logits = [math.exp(x) for x in logits]
total = sum(exp_logits)
probs = [x / total for x in exp_logits]

greedy_action = actions[probs.index(max(probs))]
sampled_action = random.choices(actions, weights=probs, k=1)[0]
sampled_prob = probs[actions.index(sampled_action)]

print("policy probabilities:", dict(zip(actions, [round(p, 3) for p in probs])))
print("greedy action:", greedy_action)
print("sampled action:", sampled_action, "log_prob:", round(math.log(sampled_prob), 3))

The worked example is small, but it shows the contract. The policy must not only choose an action; it must remember how surprising that action was under the current parameters. Without that record, the optimizer cannot say whether the action should become more or less likely after the return is known.

Practical Recipe

1. Choose the policy distribution to match the actuator: categorical for discrete actions, Gaussian or squashed Gaussian for continuous commands.
2. Log observations, sampled actions, rewards, terminations, and old log probabilities for every rollout step.
3. Keep exploration physically plausible by bounding actions and standard deviations before the command reaches the controller.
4. Evaluate with the same initial states and perturbations whenever two policy updates are compared.
5. Separate policy failure from execution failure by logging controller saturation, contact slips, timeouts, and safety stops.

A direct policy optimizer only sees the consequences of actions that were actually sampled. That makes distribution design a systems decision, not a cosmetic modeling choice. Too little entropy prevents discovery; too much entropy spends rollouts on unsafe or uninformative behavior.

For embodied agents, the action distribution also sits between learning and control. A Gaussian policy over joint velocity commands may sample a value that the safety layer clips. If training records the unclipped log probability but the robot executes the clipped command, the update credits the wrong action. The implementation must log both the policy sample and the executed command.

A robust implementation treats action sampling as part of the evidence artifact. Code Fragment 2 sketches the fields a rollout buffer needs before PPO or REINFORCE can make a valid policy-gradient update.

1. Store the observation before action sampling, not a later state estimate after the controller has moved.
2. Store the sampled action, executed action, old log probability, reward, value estimate, and termination flag in one row.
3. Record the random seed and policy version that produced the rollout.
4. Reject rollout rows where safety clipping or controller failure makes the training target ambiguous, or mark them with an explicit failure label.
5. Compare policies only when rollouts use the same reset distribution and perturbation suite.

# Define the rollout row that direct policy optimization needs.
# Store both sampled and executed actions so safety clipping is visible.
from dataclasses import dataclass, asdict

@dataclass
class RolloutRow:
    observation_id: str
    sampled_action: float
    executed_action: float
    old_log_prob: float
    reward: float
    clipped_by_safety: bool

    def as_row(self) -> dict[str, object]:
        return asdict(self)

row = RolloutRow(
    observation_id="episode_0042_step_0017",
    sampled_action=1.35,
    executed_action=1.00,
    old_log_prob=-0.42,
    reward=0.8,
    clipped_by_safety=True,
)
print(row.as_row())

When direct policy optimization fails, first inspect the action distribution before blaming the optimizer. Check entropy, action clipping, invalid-action masking, controller saturation, and whether high-return episodes came from meaningful exploration or lucky resets. Then rerun a small perturbation panel with fixed seeds so the policy change is compared against the same embodied conditions.