Policy Gradient and Actor–Critic · Study

English translation

Execution conditions

Required packages: gymnasium, stable-baselines3, and PyTorch.
SB3 policy: MlpPolicy.
Gymnasium step: 5 return values.

If the required package is missing, run %pip install gymnasium==1.3.0 stable-baselines3==2.9.2 in a Colab or Jupyter code cell.

Topics covered

Examine direct updates to a stochastic policy, then investigate how a critic changes learning over long episodes.

Update the policy directly

REINFORCE weights the log probability of the selected action by the return. Subtracting a baseline does not change the reward function itself.

The critic's role

A2C uses a critic to estimate future values, computes advantages, and updates the actor. Rollout length changes the balance between actual rewards and bootstrapping.

Worked example

Queue, softmax policy 1→16→3, $\gamma=.95$ , Adam .005, and 30 episodes. Use a simple baseline that subtracts the mean return for each episode.

"""Synthetic service queue. An action chooses capacity, including its energy cost."""
import numpy as np

class Queue:
    def reset(self, seed=0):
        self.rng = np.random.default_rng(seed)
        self.q, self.t = 0, 0
        return self.q

    def step(self, action):
        capacity = (0, 2, 4)[int(action)]
        arriving = int(self.rng.poisson(2))
        raw = max(0, self.q - capacity) + arriving
        dropped = max(0, raw - 20)
        self.q = min(20, raw)
        self.t += 1
        reward = -(self.q + .2 * capacity**2 + 10 * dropped)
        return self.q, reward, False, self.t >= 100, {'queue': self.q, 'capacity': capacity, 'dropped': dropped}

def evaluate(policy, seeds=range(100, 105)):
    rows = []
    for seed in seeds:
        env = Queue(); q = env.reset(seed)
        total = waiting = energy = dropped = 0
        while True:
            q, r, terminated, truncated, info = env.step(policy(q))
            total += r; waiting += q; energy += info['capacity']**2; dropped += info['dropped']
            if terminated or truncated: break
        rows.append([total, waiting / env.t, energy / env.t, dropped])
    return np.array(rows)

policies = {'constant': lambda q: 1, 'threshold': lambda q: 0 if q == 0 else (1 if q < 5 else 2)}
for name, policy in policies.items():
    print(name, evaluate(policy).mean(axis=0))


import torch
from torch import nn
from torch.distributions import Categorical
def fit_reinforce(seed=0,baseline=True):
    torch.manual_seed(seed)
    model=nn.Sequential(nn.Linear(1,16),nn.Tanh(),nn.Linear(16,3))
    optimizer=torch.optim.Adam(model.parameters(),lr=.005);history=[]
    for episode in range(30):
        env=Queue();q=env.reset(seed*1000+episode);logs=[];rewards=[]
        for k in range(100):
            distribution=Categorical(logits=model(torch.tensor([[q/20]],dtype=torch.float32)))
            a=distribution.sample();q,r,te,tr,_=env.step(int(a.item()))
            logs.append(distribution.log_prob(a).squeeze());rewards.append(r/20)
            if te or tr:break
        returns=[];g=0
        for r in rewards[::-1]:g=r+.95*g;returns.append(g)
        returns=torch.tensor(returns[::-1],dtype=torch.float32)
        advantage=returns-returns.mean() if baseline else returns
        loss=-(torch.stack(logs)*advantage).mean()
        optimizer.zero_grad();loss.backward();optimizer.step();history.append(sum(rewards)*20)
    def policy(q):
        with torch.no_grad():return int(model(torch.tensor([[q/20]],dtype=torch.float32)).argmax())
    return policy,history
policy,history=fit_reinforce()
print(evaluate(policy).mean(0))

constant [-921.2      8.052    4.       3.6  ]
threshold [-387.4     2.77    5.52    0.  ]
[-531.4      2.114   16.       0.   ]

Official documentation and data

SB3 · A2C

Reference example

View more examples