Policy Parameterization and Objectives

A parameterized policy is a function $\pi_\theta(a \mid s)$ that maps a state $s$ to a probability distribution over actions $a$, where $\theta \in \mathbb{R}^d$ is a real-valued parameter vector — for instance the weights and biases of a neural network, or the coefficients of a linear feature map. This is a generalization of the tabular policies used in Markov Decision Processes: instead of storing a separate probability for every state-action pair, we store a compact, differentiable rule that computes those probabilities on demand, generalizing across states that share structure. The word "stochastic" matters: $\pi_\theta$ outputs a distribution, not a single deterministic action, both because randomized exploration is often necessary and because in many domains (e.g. games with hidden information) the optimal policy genuinely is stochastic. Two canonical parameterizations recur throughout the field. For discrete action spaces, we compute a real-valued preference $h(s,a,\theta)$ for each action — typically the output of a neural network with one number per action — and convert preferences to probabilities with the softmax:

For continuous action spaces, we instead parameterize the moments of a distribution directly, most commonly a Gaussian:

where $\mu_\theta(s)$ and $\sigma_\theta(s)$ are themselves outputs of a function approximator. This chapter foreshadows, but does not fully develop, the machinery of Continuous Action Spaces and Gaussian Policies — the key point for now is simply that parameterizing a distribution's parameters rather than an action directly is what makes the policy differentiable and hence trainable by gradient methods. The objective $J(\theta)$ is the scalar we are trying to maximize by choosing $\theta$: the expected return obtained by following $\pi_\theta$. Its precise definition depends on the task structure. In the episodic case, with a designated start-state distribution, we define

the expected discounted return from the start. In the continuing case, where there is no terminal state and episodes never end, discounted return can diverge or lose meaning, so the standard objective becomes the average reward rate:

where $d_{\pi_\theta}$ is the stationary state distribution induced by following $\pi_\theta$ forever. Both formulations share the same essential shape: $J(\theta)$ is an expectation, under trajectories generated by $\pi_\theta$, of accumulated reward. Policy-based reinforcement learning is precisely the problem of finding $\theta^* = \arg\max_\theta J(\theta)$, or more realistically, of climbing $J(\theta)$'s gradient as far as we can, since $\theta^*$ is rarely reachable in closed form.

Turning "maximize $J(\theta)$" into an algorithm requires answering three separate questions, and the clean separation between them is exactly what makes the framework modular and reusable across problem types.

Question 1: How do we represent $\pi_\theta$? We choose a functional form (linear-softmax, neural softmax, Gaussian head) and initialize $\theta$, typically at random or near zero so the initial policy is close to uniform/exploratory.

Question 2: How do we evaluate $J(\theta)$, or at least estimate it? We cannot compute the expectation over all possible trajectories in closed form for any nontrivial MDP. Instead, we sample trajectories by rolling the current policy out in the environment (or a simulator) and use the empirical return as an unbiased (Monte Carlo) estimate of $J(\theta)$.

Question 3: How do we improve $\theta$? This chapter sets up the objective but defers the differentiation itself to the Policy-Gradient Theorem; here we simply establish the loop structure that every policy-gradient method shares, so that later chapters can slot in the actual gradient estimator. ``` PolicyOptimizationLoop(theta_init, num_iterations, episodes_per_batch): theta <- theta_init for iteration = 1 to num_iterations: batch_trajectories <- empty list

# --- Question 2: estimate J(theta) by sampling rollouts --- for episode = 1 to episodes_per_batch: trajectory <- empty list s <- sample initial state from d0 while s is not terminal (and episode not truncated): a <- sample action from pi_theta( . | s ) # stochastic policy s', r <- step environment(s, a) trajectory.append( (s, a, r) ) s <- s' batch_trajectories.append(trajectory) J_hat <- average total (discounted) return over batch_trajectories

# J_hat is an unbiased Monte Carlo estimate of J(theta)

# --- Question 3: improve theta using J_hat / trajectories --- grad_estimate <- EstimatePolicyGradient(batch_trajectories, theta)

# <- defined by the Policy-Gradient Theorem, deferred here theta <- theta + alpha grad_estimate # gradient ASCENT on J(theta) return theta ``` Notice the loop's structure: it alternates between acting (sampling actions from the current stochastic policy, which supplies exploration for free) and learning* (nudging $\theta$ using information gathered from those actions' outcomes). Nothing in this loop depends on knowing the environment's transition dynamics $p(s',r\mid s,a)$ — only the ability to sample from them. This is what makes policy-based methods model-free, and it is the same property that lets the same skeleton support both discrete-action softmax policies and continuous-action Gaussian policies, since neither the sampling step nor the update step depends on which parameterization is used.

Representational complexity. If $\pi_\theta$ is a linear-softmax policy over $|\mathcal{A}|$ actions with a $k$-dimensional feature vector, then evaluating $\pi_\theta(a\mid s)$ for all actions costs $O(k|\mathcal{A}|)$ time (one dot product per action) and $O(k|\mathcal{A}|)$ space for $\theta$. For a neural policy with $L$ layers and width $w$, a forward pass costs $O(Lw^2)$ time and the parameter count $d = O(Lw^2)$ dominates storage — independent of the size of the state space, which is the entire point: a tabular policy over $|\mathcal{S}|$ states would need $\Omega(|\mathcal{S}||\mathcal{A}|)$ parameters, intractable for continuous or combinatorially large state spaces, whereas a parameterized policy's cost depends only on the chosen function class.

Sampling complexity of estimating $J(\theta)$. Because $J(\theta) = \mathbb{E}_{\pi_\theta}[G]$ has no closed form for general MDPs, we estimate it by Monte Carlo: draw $N$ trajectories, average their returns. This estimator is unbiased, since each trajectory's return is drawn from exactly the distribution induced by $\pi_\theta$, and by the law of large numbers its variance shrinks as $O(1/N)$, so its standard error shrinks as $O(1/\sqrt N)$ — the same $1/\sqrt N$ rate as any Monte Carlo estimator, regardless of state or action space size. This is a genuine tradeoff worth stating honestly: policy-based methods pay for their model-freedom and scalability with statistically noisy objective and gradient estimates, one motivation for the variance-reduction techniques (baselines, actor-critic) developed later.

Correctness of the objective's definition. A subtlety worth checking carefully: is $J(\theta)$ actually a well-defined, finite number for all $\theta$? In the episodic case with a proper policy (one that reaches termination with probability 1) and bounded per-step reward, $G_0$ is a finite random variable with bounded expectation, so $J(\theta)$ is well-defined. In the continuing case, the average-reward limit exists whenever the Markov chain induced by $\pi_\theta$ is ergodic (irreducible and aperiodic), guaranteeing a unique stationary distribution $d_{\pi_\theta}$ — this is precisely the technical condition under which the average-reward formula above is valid; without ergodicity, $J(\theta)$ may depend on the start state or fail to converge, which is why textbook treatments typically restrict attention to ergodic MDPs when using this objective.

Non-concavity, honestly stated. Even once $J(\theta)$ is well defined, it is essential to be upfront that $J$ is, in general, a highly non-concave function of $\theta$ for any nonlinear parameterization (in particular, any neural-network policy). Gradient ascent on $J(\theta)$ is therefore only guaranteed to find a local maximum (or, in the noisy stochastic setting, a stationary point of the expected objective), not the global optimum — a real limitation shared with the rest of deep learning, and not something this framework circumvents. What can be shown rigorously (and is the content of the Policy-Gradient Theorem) is that the direction $\nabla_\theta J(\theta)$ is well-defined and computable from sampled trajectories alone, without needing a model of the environment — that theorem is what turns "climb $J(\theta)$" from a slogan into an algorithm.

Consider a tiny discrete-action MDP: a 1-D "gridworld corridor" with states $\{0,1,2\}$, where state $2$ is terminal (goal, reward $+1$), and state $0$ gives reward $-1$ per step to discourage lingering. Two actions are available in state $1$: `left` and `right`. We use a linear-softmax policy with a single scalar feature per action: $h(1,\text{right},\theta) = \theta_r$, $h(1,\text{left},\theta) = \theta_l$, so

Suppose we initialize $\theta_r = \theta_l = 0$, giving $\pi_\theta(\text{right}\mid 1) = 0.5$: uniform, maximally exploratory, as expected at initialization.

Two things this trace makes concrete. First, $J(\theta)$ at iteration 1 is genuinely estimated from noisy samples — the two episodes shown gave wildly different returns ($+1$ and $-4$) from the same $\theta$, purely because the stochastic policy sampled different actions; this is exactly the Monte-Carlo variance discussed above. Second, as $\theta_r$ is pushed up relative to $\theta_l$ (by whatever gradient rule the later chapter supplies), the softmax smoothly moves probability mass toward the better action without ever violating $\sum_a \pi_\theta(a\mid s) = 1$ — a guarantee that comes for free from the softmax's construction, and would have to be re-derived from scratch if we instead tried to parameterize probabilities directly and clip them into $[0,1]$.

The idea of directly parameterizing and optimizing a policy, rather than deriving one from value functions, has roots that predate deep learning and connects reinforcement learning to statistics, control theory, and stochastic approximation.

Nearly every large-scale application of modern reinforcement learning rests on exactly the framework introduced here. Robotic control systems use Gaussian-parameterized continuous-action policies to output joint torques or motor commands, because tabular or discretized action representations are infeasible for high-dimensional actuators. Game-playing agents in complex action spaces (e.g. real-time strategy or continuous-control benchmarks) use softmax or Gaussian policies parameterized by deep networks, trained by maximizing sample-estimated $J(\theta)$. Reinforcement learning from human feedback (RLHF), used to align large language models, casts token generation itself as a stochastic policy $\pi_\theta(\text{token}\mid \text{context})$ and fine-tunes $\theta$ to maximize an expected-reward objective structurally identical to $J(\theta)$ here. Simulation-based training pipelines in autonomous driving and industrial control likewise define an explicit reward signal and a parameterized controller, then run the same sample-estimate-and-improve loop. Any system where "what to do" must be learned directly, rather than derived from an intermediate value table, depends on the concepts in this article — the policy representation and the expected-return objective are the shared substrate underneath the entire actor-critic and policy-gradient literature, including the theorem that differentiates $J(\theta)$ and the Gaussian policies used for continuous control.

Ask the tutor: How does the Policy-Gradient Theorem let us compute $\nabla_\theta J(\theta)$ using only sampled trajectories, without ever differentiating through the environment's unknown transition dynamics?