Deep Reinforcement Learning with CartPole: Policy Gradient and Actor-Critic Methods Explained

Deep reinforcement learning with the CartPole environment teaches agents to balance a pole using neural network policies trained via Policy Gradient (REINFORCE) or Actor-Critic algorithms, implemented in both TensorFlow and PyTorch.

The CartPole benchmark from OpenAI Gym serves as the canonical introduction to deep RL. Microsoft's AI-For-Beginners curriculum demonstrates how neural networks learn control policies directly from environment interaction—without hand-crafted features or model-based planning. This guide examines the two fundamental architectures used: Policy Gradient and Actor-Critic, with direct references to the repository's implementation.

The CartPole Environment: State, Action, and Reward

The environment provides a minimal yet complete testbed for continuous control. As defined in CartPole-RL-TF.ipynb (lines 42–48), the problem structure is:

• State observation – a 4-dimensional vector: [cart_position, cart_velocity, pole_angle, pole_angular_velocity]
• Action space – discrete binary choice: 0 (push left) or 1 (push right)
• Reward signal – +1 for each timestep the pole remains upright
• Termination – episode ends when pole angle exceeds ±12° or cart position exceeds ±2.4 (lines 88–91)

The agent's goal is to maximize cumulative reward, typically reaching the "solved" threshold of 200 timesteps per episode.

Policy Gradient (REINFORCE): Direct Policy Optimization

The Policy Gradient method learns a stochastic policy π(a|s) that maps states to action probabilities. The AI-For-Beginners implementation follows the classic REINFORCE algorithm.

Network Architecture

Both framework implementations use a shallow feed-forward architecture:

• Input layer: 4 neurons (state dimensions)
• Hidden layer: 128 neurons with ReLU activation
• Output layer: 2 neurons with softmax (action probabilities)

TensorFlow version (CartPole-RL-TF.ipynb, lines 28–31):

model = keras.Sequential([
    keras.layers.Dense(128, activation="relu", input_shape=(4,)),
    keras.layers.Dense(2, activation="softmax")
])
model.compile(optimizer=keras.optimizers.Adam(1e-2), loss='categorical_crossentropy')

PyTorch equivalent (CartPole-RL-PyTorch.ipynb, lines 13–18) defines the same structure using nn.Linear and F.softmax.

Episode Collection and Training

The training loop collects experience by running the environment:

1. Sample action from the current policy distribution
2. Execute action, observe reward and next state
3. Store trajectory: (state, action, log_probability, reward)

TensorFlow implementation (lines 49–55):

def run_episode(max_steps=10000):
    states, actions, probs, rewards = [], [], [], []
    state = env.reset()
    for _ in range(max_steps):
        probs_vec = model(np.expand_dims(state,0))[0].numpy()
        action = np.random.choice(num_actions, p=probs_vec)
        next_state, reward, done, _ = env.step(action)
        states.append(state); actions.append(action)
        probs.append(probs_vec); rewards.append(reward)
        state = next_state
        if done: break
    return np.array(states), np.array(actions), \
           np.array(probs), np.array(rewards)

Discounted Returns and Gradient Update

The algorithm computes discounted returns G_t to credit past actions for future rewards:

def discounted_rewards(rews, gamma=0.99):
    dr = np.zeros_like(rews, dtype=np.float32)
    running = 0
    for t in reversed(range(len(rews))):
        running = rews[t] + gamma * running
        dr[t] = running
    return (dr - dr.mean()) / (dr.std() + 1e-8)  # normalize

The loss function maximizes expected return via gradient ascent on log π(a|s) · G_t. In practice, this becomes minimization of -log π(a|s) · G_t using Adam (TensorFlow lines 70–78; PyTorch lines 22–27).

The policy typically converges to the 200-reward threshold within 300 episodes.

Actor-Critic: Reducing Variance with Value Estimation

The Actor-Critic architecture addresses Policy Gradient's high variance by introducing a learned value function V(s) as a baseline. This creates two interacting components:

• Actor – the policy network that selects actions
• Critic – the value network that estimates expected future returns

Dual-Head Architecture

TensorFlow implements this with shared hidden layers feeding two outputs (CartPole-RL-TF.ipynb, lines 66–84). PyTorch uses separate classes (CartPole-RL-PyTorch.ipynb, lines 106–124):

class Actor(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(state_dim, 128)
        self.fc2 = nn.Linear(128, action_dim)
    def forward(self, x):
        x = F.relu(self.fc1(x))
        probs = F.softmax(self.fc2(x), dim=-1)
        return torch.distributions.Categorical(probs)

class Critic(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(state_dim, 128)
        self.fc2 = nn.Linear(128, 1)
    def forward(self, x):
        x = F.relu(self.fc1(x))
        return self.fc2(x)

Advantage-Based Updates

The key innovation is the advantage function: A_t = G_t - V(s_t)

• Positive advantage → action was better than expected → strengthen
• Negative advantage → action was worse than expected → weaken

Training loop (PyTorch, lines 104–113):

for episode in range(200):
    log_probs, values, rewards, masks, next_state = run_episode()
    returns = compute_returns(next_state, rewards, masks)
    returns = torch.tensor(returns)
    values = torch.cat(values).squeeze()
    advantage = returns - values  # lower variance than returns alone

    actor_loss = -(torch.stack(log_probs) * advantage.detach()).mean()
    critic_loss = advantage.pow(2).mean()  # minimize value error

    opt_a.zero_grad(); opt_c.zero_grad()
    actor_loss.backward(); critic_loss.backward()
    opt_a.step(); opt_c.step()

The Actor-Critic variant converges faster than vanilla Policy Gradient because the critic provides a lower-variance baseline for policy updates.

Key Implementation Files

File	Location	Description
CartPole-RL-TF.ipynb	lessons/6-Other/22-DeepRL/	TensorFlow implementations of Policy Gradient and Actor-Critic
CartPole-RL-PyTorch.ipynb	lessons/6-Other/22-DeepRL/	PyTorch versions with separate Actor/Critic classes
MountainCar.ipynb	lessons/6-Other/22-DeepRL/lab/	Extended RL lab applying same concepts to continuous control
notebook.ipynb	lessons/6-Other/22-DeepRL/	Consolidated theory-code walkthrough

Summary

• CartPole provides a 4-dimensional state space and binary action space, ideal for demonstrating deep RL fundamentals
• Policy Gradient (REINFORCE) directly optimizes a neural policy using Monte Carlo returns, implemented with 128-unit hidden layers in both TensorFlow and PyTorch
• Actor-Critic reduces gradient variance by learning a value function baseline, achieving faster convergence through advantage-based updates
• Both architectures solve the environment at ~200 reward threshold, with Actor-Critic typically requiring fewer episodes
• The AI-For-Beginners notebooks provide complete, runnable implementations with detailed annotations for educational use

Frequently Asked Questions

What makes CartPole a good environment for learning deep RL?

CartPole offers the simplest possible continuous control problem: low-dimensional state (4 values), discrete actions (left/right), immediate rewards, and fast episode execution. These properties enable rapid experimentation and clear visualization of policy improvements, making abstract RL concepts concrete before tackling complex environments like Atari games or robotics.

How does the Actor-Critic advantage function reduce variance?

The advantage A_t = G_t - V(s_t) measures how much better an action was than the average action in that state. By subtracting the learned value baseline, the gradient update depends only on relative performance rather than absolute return magnitude. This removes common-mode noise from reward scaling and environment stochasticity, as implemented in the repository's advantage.pow(2).mean() critic loss.

Why normalize discounted returns in Policy Gradient?

Return normalization (subtract mean, divide by standard deviation) ensures gradient magnitudes remain stable across episodes with varying lengths and reward totals. Without normalization, long episodes with high cumulative returns would dominate the gradient, causing unstable updates. The repository's discounted_rewards() function applies this standardization with a small epsilon for numerical stability.

Can these methods scale to more complex environments?

The same architectural principles apply to larger problems, though significant modifications become necessary: deeper networks (e.g., CNNs for visual states), experience replay buffers, target networks for the critic (as in DQN), and more sophisticated advantage estimators (GAE, PPO-style clipped objectives). The CartPole implementations in AI-For-Beginners deliberately omit these complications to focus on core algorithmic concepts.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}