How to Implement Deep Reinforcement Learning with PyTorch Using CartPole-RL-PyTorch.ipynb

The CartPole-RL-PyTorch.ipynb notebook provides a complete implementation of policy-gradient and actor-critic algorithms using PyTorch to solve the CartPole-v1 environment, including trajectory collection, discounted reward calculation, and network training loops.

The microsoft/AI-For-Beginners repository contains a hands-on tutorial for deep reinforcement learning with PyTorch in lessons/6-Other/22-DeepRL/CartPole-RL-PyTorch.ipynb. This notebook demonstrates how to train neural network agents using two classic algorithms: a simple policy-gradient method and a more advanced actor-critic architecture. Both approaches use the OpenAI Gym CartPole-v1 environment, where the agent learns to balance a pole by applying forces to a cart.

Setting Up the CartPole Environment

Before implementing the algorithms, the notebook initializes the simulation environment and inspects its properties. The CartPole-v1 environment provides a 4-dimensional continuous observation space (cart position, cart velocity, pole angle, pole angular velocity) and a 2-dimensional discrete action space (push left or right).

import gym
import numpy as np
import torch
import matplotlib.pyplot as plt

env = gym.make("CartPole-v1")
print(f"Action space: {env.action_space}")
print(f"Observation space: {env.observation_space}")

Running this setup confirms the state dimensionality (4) and action count (2), which determine the input and output sizes of your neural networks.

Policy-Gradient Implementation

The notebook first implements a policy-gradient (REINFORCE) algorithm using a feed-forward neural network that directly maps states to action probabilities.

Building the Policy Network

The policy network uses torch.nn.Sequential to define a simple architecture with one hidden layer. The input layer accepts 4 state dimensions, the hidden layer uses 128 ReLU units, and the output layer produces 2 softmax probabilities corresponding to the left and right actions.

model = torch.nn.Sequential(
    torch.nn.Linear(4, 128, bias=False),
    torch.nn.ReLU(),
    torch.nn.Linear(128, 2, bias=False),
    torch.nn.Softmax(dim=1)
)

Collecting Trajectories with run_episode()

The run_episode() function executes a full episode by sampling actions from the current policy distribution, storing the resulting states, actions, probabilities, and rewards for later gradient computation.

def run_episode(max_steps=10000, render=False):
    states, actions, probs, rewards = [], [], [], []
    state = env.reset()
    for _ in range(max_steps):
        if render:
            env.render()
        action_probs = model(torch.from_numpy(state[None, :]))[0]
        action = np.random.choice(2, p=action_probs.detach().numpy())
        next_state, reward, done, _ = env.step(action)
        states.append(state)
        actions.append(action)
        probs.append(action_probs.detach().numpy())
        rewards.append(reward)
        if done:
            break
        state = next_state
    return np.vstack(states), np.vstack(actions), np.vstack(probs), np.vstack(rewards)

This function returns complete trajectories necessary for Monte Carlo policy gradient updates.

Computing Discounted Returns

The discounted_rewards() function calculates the return for each timestep by summing future rewards with a decay factor (gamma=0.99). It optionally normalizes the returns to reduce variance in gradient estimates.

def discounted_rewards(rewards, gamma=0.99, normalize=True):
    ret, s = [], 0
    for r in rewards[::-1]:
        s = r + gamma * s
        ret.insert(0, s)
    if normalize:
        eps = 1e-4
        ret = (np.array(ret) - np.mean(ret)) / (np.std(ret) + eps)
    return ret

Normalization stabilizes training by ensuring reward scales remain consistent across episodes.

Training Loop Implementation

The training loop uses Adam optimization with a learning rate of 0.01. For each episode, it computes policy gradients by comparing one-hot encoded actions against predicted probabilities, weighted by discounted returns.

optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
alpha = 1e-4
history = []

for epoch in range(300):
    states, actions, probs, rewards = run_episode()
    one_hot = np.eye(2)[actions.T][0]
    grads = one_hot - probs
    dr = discounted_rewards(rewards)
    grads *= dr
    target = alpha * np.vstack([grads]) + probs
    
    loss = -torch.mean(torch.log(model(torch.from_numpy(states))) * torch.from_numpy(target))
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    history.append(np.sum(rewards))
    if epoch % 100 == 0:
        print(f"Epoch {epoch}: Total Reward = {np.sum(rewards)}")

plt.plot(history)

This loop demonstrates how to manually construct policy-gradient loss functions without relying on high-level RL libraries.

Actor-Critic Architecture

The notebook extends the approach with an actor-critic method, which reduces variance by learning a value function alongside the policy.

Defining Actor and Critic Networks

Two separate networks share the same input dimension but output different predictions. The Actor outputs a categorical distribution over actions, while the Critic estimates the state-value V(s). Both networks use three hidden layers (128 → 256 → output) and automatically utilize CUDA if available.

class Actor(torch.nn.Module):
    def __init__(self, s_dim, a_dim):
        super().__init__()
        self.net = torch.nn.Sequential(
            torch.nn.Linear(s_dim, 128),
            torch.nn.ReLU(),
            torch.nn.Linear(128, 256),
            torch.nn.ReLU(),
            torch.nn.Linear(256, a_dim)
        )
    
    def forward(self, state):
        logits = self.net(state)
        return torch.distributions.Categorical(
            torch.nn.functional.softmax(logits, dim=-1)
        )

class Critic(torch.nn.Module):
    def __init__(self, s_dim):
        super().__init__()
        self.net = torch.nn.Sequential(
            torch.nn.Linear(s_dim, 128),
            torch.nn.ReLU(),
            torch.nn.Linear(128, 256),
            torch.nn.ReLU(),
            torch.nn.Linear(256, 1)
        )
    
    def forward(self, state):
        return self.net(state)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
actor = Actor(4, 2).to(device)
critic = Critic(4).to(device)

Training with Advantage Estimation

The actor-critic training loop computes the advantage (returns minus predicted values) to update both networks. The actor loss maximizes log-probabilities weighted by advantage, while the critic loss minimizes the mean squared error of value predictions.

def run_episode(actor, critic, n_iters):
    log_probs, values, rewards, masks = [], [], [], []
    state = env.reset()
    
    for _ in range(n_iters):
        state_tensor = torch.from_numpy(state).float().to(device)
        dist = actor(state_tensor)
        value = critic(state_tensor)
        
        action = dist.sample()
        next_state, reward, done, _ = env.step(action.cpu().numpy())
        
        log_prob = dist.log_prob(action)
        log_probs.append(log_prob)
        values.append(value)
        rewards.append(torch.tensor([reward], device=device))
        masks.append(torch.tensor([1 - done], device=device))
        
        state = next_state
        
        if done:
            break
    
    next_value = critic(torch.from_numpy(next_state).float().to(device))
    returns = discounted_rewards(next_value, rewards, masks)
    
    log_probs = torch.stack(log_probs)
    returns = torch.cat(returns).detach()
    values = torch.cat(values)
    
    advantage = returns - values
    
    actor_loss = -(log_probs * advantage.detach()).mean()
    critic_loss = advantage.pow(2).mean()
    
    optimizer_actor.zero_grad()
    optimizer_critic.zero_grad()
    actor_loss.backward()
    critic_loss.backward()
    optimizer_actor.step()
    optimizer_critic.step()

As implemented in CartPole-RL-PyTorch.ipynb, this approach typically converges faster than vanilla policy-gradient due to the reduced variance from baseline subtraction.

Visualization and Evaluation

After training, the notebook visualizes learning progress by plotting the reward history using matplotlib. You can evaluate the final policy by running run_episode(render=True), which renders the CartPole simulation to observe the trained agent's behavior.

# Plot training progress

plt.plot(history)
plt.xlabel("Episode")
plt.ylabel("Total Reward")
plt.title("Policy Gradient Training Progress")

# Visualize final policy

run_episode(render=True)

Summary

• CartPole-RL-PyTorch.ipynb implements two deep reinforcement learning with PyTorch algorithms: policy-gradient (REINFORCE) and actor-critic.
• The policy network maps 4-dimensional states to 2 action probabilities using a 128-unit hidden layer.
• Discounted returns are computed with gamma=0.99 and normalized to stabilize gradient estimates.
• The actor-critic architecture uses separate 128→256 layer networks for policy and value estimation, trained with advantage-based loss functions.
• Both methods use Adam optimizers and can run on CPU or CUDA devices.

Frequently Asked Questions

What is the difference between the policy-gradient and actor-critic implementations in the notebook?

The policy-gradient method updates the network using only the discounted returns as weights for log-probabilities, while the actor-critic method uses a learned value function (the critic) to compute advantages (returns minus values). This baseline subtraction in actor-critic reduces variance and typically enables faster, more stable learning compared to the simpler policy-gradient approach.

Why does the notebook normalize discounted rewards?

Normalization ensures that the gradient updates remain scale-invariant across episodes. Without normalization, a sequence of high-reward episodes would produce massive gradients, while low-reward episodes would produce vanishing gradients. The discounted_rewards() function subtracts the mean and divides by the standard deviation (with epsilon=1e-4 for numerical stability) to maintain consistent optimization dynamics.

What hyperparameters does CartPole-RL-PyTorch.ipynb use?

The notebook uses an Adam optimizer with a learning rate of 0.01 for the policy-gradient method, a discount factor (gamma) of 0.99 for computing returns, and a small scaling factor (alpha=1e-4) for gradient updates. The actor-critic networks use three hidden layers with 128 and 256 units respectively, though specific learning rates for the actor-critic optimizers can be configured based on convergence behavior.

Can I run this implementation without a GPU?

Yes. The code automatically detects device availability using torch.device("cuda" if torch.cuda.is_available() else "cpu") and moves networks to the appropriate device. The CartPole environment is lightweight enough to train efficiently on CPU, typically converging within minutes for the policy-gradient method and even faster for actor-critic.

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