9 The Calculus of Decisions: Policy Gradient Methods

This chapter introduces policy gradient methods as a way to optimize decision-making policies directly, rather than first learning action values and deriving behavior indirectly. The motivation comes from problems like data center cooling, where actions are continuous, high-dimensional, and tightly coupled across many interacting zones. In such settings, value-based methods like DQN become awkward or infeasible because discretizing the action space destroys precision and causes an explosion in the number of possible action combinations. Policy gradients avoid this bottleneck by learning a parameterized policy that maps states straight to actions, while naturally supporting stochastic behavior that can improve exploration and robustness.

The mathematical core is the policy gradient theorem, which shows how to compute the gradient of expected return without knowing the environment dynamics. Using the log-derivative trick, the update depends on the gradient of the policy’s log-probability for the actions that were taken, multiplied by the return from those actions onward. This leads to the REINFORCE algorithm, a Monte Carlo method that samples trajectories, computes return-to-go, and increases the probability of actions that led to good outcomes while decreasing the probability of poor ones. The chapter also explains how Gaussian policies handle continuous actions and how softmax policies handle discrete actions, while emphasizing the role of the reparameterization trick and the importance of correct log-probability calculations when actions are squashed into bounded ranges.

Because vanilla REINFORCE can have high variance, the chapter adds a learned baseline, usually a value network, to form advantages and make learning signals more stable and informative. Entropy regularization is then used to prevent premature collapse to overly deterministic behavior, keeping exploration alive during training. To handle the multi-zone structure of the cooling problem, the chapter introduces attention and Transformer architectures so that each zone can condition its decisions on relevant relationships with other zones. The resulting Transformer-based actor-critic agent is trained on a realistic data center cooling simulation, where it learns to balance energy use against thermal safety and outperforms a hand-crafted PI controller, demonstrating the practical strength of policy gradients for complex continuous-control tasks.

Value-Based vs. Policy-Based Reinforcement Learning.

The Policy Gradient Objective.

Mathematical proof of Policy Gradient Theorem.

The REINFORCE algorithm core steps.

Variance Reduction with Baselines.

The Attention Mechanism — Query, Key, Value.

Multi head attention architecture

Transformer architecture.

Cooling optimization problem.

REINFORCE Training Progress — Data Center Cooling

RL Policy vs. Heuristic Controller — Box Plot Comparison

• Value-based methods like DQN require discrete, enumerable action spaces. When actions are continuous (fan speeds, valve openings, motor torques) or high-dimensional (multiple agents or zones), the action space cannot be discretized without losing precision or creating combinatorial explosion.
• Policy gradient methods optimize the policy directly, representing it as a parameterized probability distribution over actions. This naturally handles continuous and high-dimensional action spaces without discretization.
• The policy gradient theorem provides a way to compute ∇_θJ(θ) without knowing environment dynamics. The gradient is: E[Σₜ ∇_θ log π_θ(aₜ|sₜ) R(τ)]. The log-derivative trick eliminates the unknown transition probabilities from the gradient.
• The score function ∇_θ log π_θ(a|s) points in the direction that makes action a more likely. The return R(τ) scales this direction: positive returns reinforce actions, negative returns suppress them. This is the core mechanism of REINFORCE.
• REINFORCE collects complete trajectories, computes returns, and updates the policy to make high-return trajectories more probable. It requires no replay buffer, no target network, and no bootstrapping—just sampling and gradient ascent.
• Gaussian policies for continuous control output a mean μ and standard deviation σ per action dimension. Actions are sampled from Normal(μ, σ), and the standard deviation naturally adapts to control exploration without an external schedule.
• Vanilla REINFORCE suffers from high variance because the return R(τ) can have arbitrary scale, causing all actions in a trajectory to be reinforced or suppressed regardless of their individual contribution.
• Subtracting a baseline b(s) from the return reduces variance without introducing bias. The optimal baseline is the state-value function V(s), giving the advantage A = R - V(s), which measures how much better or worse an action was compared to expectation.
• The actor-critic architecture uses two components: an actor (policy network) updated by the advantage-weighted policy gradient, and a critic (value network) trained to predict expected returns. The critic provides the baseline that stabilizes the actor's learning.
• Entropy regularization adds a bonus for maintaining policy stochasticity, preventing premature convergence to deterministic behavior before adequate exploration.
• The attention mechanism computes weighted combinations of inputs using Query-Key-Value projections. Each element queries all others, and the resulting attention weights reflect pairwise relevance. This enables dynamic, context-dependent information routing.
• Multi-head attention runs multiple independent attention computations in parallel, each learning different relationship patterns (thermal coupling, workload correlation, pressure dynamics). Concatenating the heads captures multiple types of inter-zone interactions simultaneously.
• The Transformer encoder stacks layers of multi-head self-attention and feedforward networks with residual connections and layer normalization. Each layer enables longer-range interactions, so two layers can capture indirect zone-to-zone effects.
• Positional encoding adds sinusoidal identity signals to each zone, allowing the Transformer to learn spatial relationships (zone adjacency) alongside state similarity.
• In data center case study, the Transformer-based REINFORCE agent achieved a 53.4% improvement over a hand-tuned PI controller, learning to trade slightly higher energy cost for dramatically fewer temperature violations—a trade-off directly aligned with the reward function's priorities.
• Policy gradient methods are the foundation for modern RL algorithms (PPO, SAC, A3C) and are essential whenever actions are continuous, stochastic policies are desirable, or multi-entity coordination is required. Understanding REINFORCE deeply prepares you for these more advanced methods in subsequent chapters.

FAQ