Policy Objective Functions

Why Does This Exist?

In value-based reinforcement learning (such as Q-learning or DQN), the learning goal is implicit: minimize Bellman error until the action-value function converges to $Q^*$.

In policy-based methods, however, the agent directly optimizes a parameterized policy $\pi_\theta(a \mid s)$ via gradient ascent:

$\theta_{t+1} = \theta_t + \alpha \nabla_\theta J(\theta)$

To compute that gradient, we must define the scalar performance objective $J(\theta)$ being maximized. In real-world reinforcement learning, problems exhibit fundamentally different temporal horizons:

1. Episodic tasks that naturally terminate at a terminal goal or timeout (games, robotic grasping).
2. Continuing discounted tasks where an ongoing process is evaluated across stationary state visitations with time-preference discounting (server load balancing).
3. Continuing undiscounted tasks where an infinite-horizon process runs indefinitely and success is measured purely by steady-state throughput (chemical refining, continuous industrial heating, high-frequency trading).

Selecting the wrong objective distorts learning. Applying discounted objectives to continuing problems introduces artificial horizon truncation, while episodic objectives fail on open-ended streams.

Furthermore, defining $J(\theta)$ appears to create a major theoretical dilemma: $J(\theta)$ depends on $\theta$ both through action selection probabilities $\pi_\theta(a \mid s)$ and through the resulting state distribution $d^{\pi_\theta}(s)$. Without a special mathematical structure, computing $\nabla_\theta J(\theta)$ would require differentiating the environmental state transition dynamics $\nabla_\theta d^{\pi_\theta}(s)$—an intractable model-dependent quantity. Understanding policy objective functions reveals the central miracle of policy gradients: across all three formulations, $\nabla_\theta d^{\pi_\theta}(s)$ cancels out entirely.

Think of It Like This

How It Actually Works

The Three Canonical Mathematical Formulations

1. Episodic Start-State Objective $J_1(\theta)$

In episodic environments with a designated start state $s_0$, the objective is the expected discounted cumulative return from $s_0$:

$J_1(\theta) \doteq V^{\pi_\theta}(s_0) = \mathbb{E}_{\pi_\theta} \left[ \sum_{t=0}^T \gamma^t R_{t+1} \;\middle|\; S_0 = s_0 \right]$

If the initial state is drawn from a starting distribution $d_0(s)$, the objective generalizes to the expectation over starting states:

$J_1(\theta) \doteq \sum_{s \in \mathcal{S}} d_0(s) V^{\pi_\theta}(s)$

Under this formulation, the state visitation distribution is the unnormalized discounted on-policy state visitation frequency:

$d^{\pi_\theta}(s) \doteq \sum_{t=0}^\infty \gamma^t P(S_t = s \mid S_0 = s_0, \pi_\theta)$

2. Continuing Average Value Objective $J_{\text{avV}}(\theta)$

In continuing environments evaluated with discount factor $\gamma \in [0, 1)$, we measure the average state value weighted by the stationary distribution $d^{\pi_\theta}(s)$:

$J_{\text{avV}}(\theta) \doteq \sum_{s \in \mathcal{S}} d^{\pi_\theta}(s) V^{\pi_\theta}(s)$

Where $d^{\pi_\theta}(s) \doteq \lim_{t \to \infty} P(S_t = s \mid S_0, \pi_\theta)$ is the stationary distribution under policy $\pi_\theta$, satisfying the balance equations:

$d^{\pi_\theta}(s) = \sum_{s' \in \mathcal{S}} d^{\pi_\theta}(s') \sum_{a \in \mathcal{A}} \pi_\theta(a \mid s') P(s \mid s', a)$

3. Continuous Average Reward Objective $J_{\text{avR}}(\theta) = r(\pi_\theta)$

In continuing environments without discounting ($\gamma = 1$), accumulated returns diverge to infinity. The canonical metric is the average reward per time step (the reward rate $r(\pi)$):

$J_{\text{avR}}(\theta) \doteq r(\pi_\theta) \doteq \lim_{T \to \infty} \frac{1}{T} \sum_{t=1}^T \mathbb{E}[R_t \mid \pi_\theta] = \sum_{s \in \mathcal{S}} d^{\pi_\theta}(s) \sum_{a \in \mathcal{A}} \pi_\theta(a \mid s) \mathcal{R}_s^a$

Where $\mathcal{R}_s^a \doteq \mathbb{E}[R_{t+1} \mid S_t = s, A_t = a]$ is the expected immediate reward.

The Unifying Policy Gradient Theorem

Differentiating any of these objectives appears daunting because the state distribution $d^{\pi_\theta}(s)$ depends directly on policy parameters $\theta$:

$\nabla_\theta J(\theta) = \nabla_\theta \left( \sum_s d^{\pi_\theta}(s) V^{\pi_\theta}(s) \right) = \sum_s \left( \nabla_\theta d^{\pi_\theta}(s) V^{\pi_\theta}(s) + d^{\pi_\theta}(s) \nabla_\theta V^{\pi_\theta}(s) \right)$

Evaluating $\nabla_\theta d^{\pi_\theta}(s)$ would require knowing the environment's internal transition probabilities $P(s' \mid s, a)$, making model-free learning impossible.

The Policy Gradient Theorem (Sutton et al., 1999) establishes that for all three objectives, the state distribution gradient miraculously vanishes:

$\nabla_\theta J(\theta) \propto \sum_{s \in \mathcal{S}} d^{\pi_\theta}(s) \sum_{a \in \mathcal{A}} \nabla_\theta \pi_\theta(a \mid s) Q^{\pi_\theta}(s, a)$

Using the identity $\nabla_\theta \pi_\theta(a \mid s) = \pi_\theta(a \mid s) \nabla_\theta \log \pi_\theta(a \mid s)$, this converts directly to an expectation over trajectory samples:

$\nabla_\theta J(\theta) = \mathbb{E}_{\pi_\theta} \left[ \nabla_\theta \log \pi_\theta(A_t \mid S_t) Q^{\pi_\theta}(S_t, A_t) \right]$

This single unifying result allows REINFORCE, Actor-Critic, PPO, and TRPO to optimize policies without modeling environmental dynamics.

Worked Numerical Example

Consider a 2-state cyclic MDP with states $\mathcal{S} = \{1, 2\}$ and discount factor $\gamma = 0.8$:

• State 1: The agent chooses between:
  • Action $a_0$ (Safe): Loops $1 \to 1$ with reward $R = 1.0$.
  • Action $a_1$ (Switch): Transitions $1 \to 2$ with reward $R = 3.0$.
• State 2: Forced deterministic transition $2 \to 1$ with reward $R = 0.0$.
• Policy parameter: Let $p \doteq \pi_\theta(a_1 \mid 1) \in [0, 1]$ be the probability of taking the Switch action $a_1$ in state 1.

1. Stationary Distribution $d^\pi = [d_1, d_2]$

Transition probability matrix: $P = \begin{bmatrix} 1 - p & p \\ 1.0 & 0.0 \end{bmatrix}$

The stationary distribution satisfies $d_1 (1 - p) + d_2 (1.0) = d_1 \implies d_2 = p d_1$. Since $d_1 + d_2 = 1 \implies d_1 (1 + p) = 1$: $d_1 = \frac{1}{1 + p}, \quad d_2 = \frac{p}{1 + p}$

2. Continuous Average Reward Objective $J_{\text{avR}}(p) = r(\pi)$

Expected reward in State 1: $\bar{R}_1 = (1 - p)(1.0) + p(3.0) = 1.0 + 2p$. Expected reward in State 2: $\bar{R}_2 = 0.0$. $r(\pi) = d_1 \bar{R}_1 + d_2 \bar{R}_2 = \left( \frac{1}{1 + p} \right)(1 + 2p) + 0 = \frac{1 + 2p}{1 + p}$

3. Episodic Start-State Value $J_1(p) = V(1)$ with $\gamma = 0.8$

From State 2: $V(2) = 0 + \gamma V(1) = 0.8 V(1)$. From State 1: $Q(1, a_0) = 1.0 + \gamma V(1) = 1.0 + 0.8 V(1)$ $Q(1, a_1) = 3.0 + \gamma V(2) = 3.0 + 0.8 (0.8 V(1)) = 3.0 + 0.64 V(1)$ $V(1) = (1 - p) Q(1, a_0) + p Q(1, a_1) = (1 + 2p) + (0.8 - 0.16p) V(1)$ $V(1) [0.2 + 0.16p] = 1 + 2p \implies V(1) = \frac{1 + 2p}{0.2 + 0.16p}$

4. Continuing Average Value Objective $J_{\text{avV}}(p)$

$J_{\text{avV}}(p) = d_1 V(1) + d_2 V(2) = \left( \frac{1}{1 + p} + \frac{0.8p}{1 + p} \right) V(1) = \left( \frac{1 + 0.8p}{1 + p} \right) V(1)$

Numerical Evaluation Across Policy Values:

All three objective functions strictly increase as $p$ goes from $0 \to 1$, confirming that all three formulations agree on the optimal policy ($p^* = 1.0$) while operating under distinct mathematical units.

Code

1from typing import Dict, Tuple2
3
4def evaluate_policy_objectives(5    p: float,6    gamma: float = 0.8,7) -> Dict[str, float]:8    """Compute analytical policy objective functions for the 2-state cyclic MDP.9
10    Args:11        p: Probability of taking action a1 (Switch) in State 1.12        gamma: Discount factor for episodic and average value objectives.13
14    Returns:15        Dictionary containing stationary distribution and all 3 objective values.16    """17    if not 0.0 <= p <= 1.0:18        raise ValueError("Policy parameter p must be in [0, 1].")19
20    # 1. Stationary Distribution: d1 = 1 / (1 + p), d2 = p / (1 + p)21    d1 = 1.0 / (1.0 + p)22    d2 = p / (1.0 + p)23
24    # 2. Continuous Average Reward Objective: J_avR = (1 + 2p) / (1 + p)25    j_avr = (1.0 + 2.0 * p) / (1.0 + p)26
27    # 3. Episodic Start-State Objective: J_1 = V(1) = (1 + 2p) / (0.2 + 0.16p)28    # Derived from Bellman system: V(2) = gamma * V(1)29    denom = (1.0 - gamma) + p * gamma * (1.0 - gamma)30    v1 = (1.0 + 2.0 * p) / (denom)31    v2 = gamma * v132
33    # 4. Continuing Average Value Objective: J_avV = d1 * V(1) + d2 * V(2)34    j_avv = d1 * v1 + d2 * v235
36    return {37        "d1": d1,38        "d2": d2,39        "J_1": v1,40        "J_avV": j_avv,41        "J_avR": j_avr,42    }43
44
45def simulate_empirical_average_reward(p: float, steps: int = 100_000) -> float:46    """Empirical Monte Carlo rollout verifying continuous average reward r(pi)."""47    import random48
49    random.seed(3)50    state = 151    total_reward = 0.052
53    for _ in range(steps):54        if state == 1:55            action = 1 if random.random() < p else 056            if action == 0:57                reward = 1.058                state = 159            else:60                reward = 3.061                state = 262        else:  # state == 263            reward = 0.064            state = 165        total_reward += reward66
67    return total_reward / steps68
69
70if __name__ == "__main__":71    test_policies = [0.0, 0.5, 1.0]72
73    print("Policy (p) | d(s) = [d1, d2]   | J_1 (Start) | J_avV (Avg Val) | J_avR (Avg Reward)")74    print("-" * 76)75
76    for p in test_policies:77        res = evaluate_policy_objectives(p, gamma=0.8)78        print(79            f"  p = {p:3.1f}  | [{res['d1']:.4f}, {res['d2']:.4f}] |   {res['J_1']:7.4f}   |"80            f"     {res['J_avV']:7.4f}     |     {res['J_avR']:7.4f}"81        )82
83    # Validate empirical simulation against analytical average reward for p=0.584    empirical_r = simulate_empirical_average_reward(p=0.5, steps=100_000)85    print(f"\nEmpirical Monte Carlo Average Reward (p=0.5): {empirical_r:.4f}")86    print(f"Analytical Continuous Average Reward (p=0.5): 1.3333")

1# Expected Output:2# Policy (p) | d(s) = [d1, d2]   | J_1 (Start) | J_avV (Avg Val) | J_avR (Avg Reward)3# ----------------------------------------------------------------------------4#   p = 0.0  | [1.0000, 0.0000] |    5.0000   |      5.0000     |      1.00005#   p = 0.5  | [0.6667, 0.3333] |    7.1429   |      6.6667     |      1.33336#   p = 1.0  | [0.5000, 0.5000] |    8.3333   |      7.5000     |      1.50007# 8# Empirical Monte Carlo Average Reward (p=0.5): 1.33339# Analytical Continuous Average Reward (p=0.5): 1.3333

Watch Out For

The Quick Version

• Three Canonical Metrics: Performance is formalized via Episodic Start-State Value ($J_1 = V(s_0)$), Continuing Average Value ($J_{\text{avV}} = \sum d^\pi(s) V(s)$), or Continuous Average Reward ($J_{\text{avR}} = r(\pi)$).
• The Policy Gradient Miracle: The Policy Gradient Theorem proves that the gradient of the state distribution $\nabla_\theta d^{\pi_\theta}(s)$ cancels out across all three formulations, enabling model-free optimization.
• Match Horizon to Objective: Use episodic objectives for tasks with designated start and termination states; use average reward formulations for continuing infinite-horizon industrial processes.
• Avoid False Discounting: Discounting in continuing tasks introduces artificial horizon truncation; true non-terminating control requires differential returns anchored to the average reward rate.

What to Read Next

• Policy Gradient Methods
• The REINFORCE Algorithm
• Advantage Actor-Critic (A2C)
• Phasic Policy Gradient (PPG)