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

> Learn deep reinforcement learning with PyTorch using the CartPole-RL-PyTorch notebook. Explore policy-gradient and actor-critic algorithms with complete implementations and training loops.

- Repository: [Microsoft/AI-For-Beginners](https://github.com/microsoft/AI-For-Beginners)
- Tags: tutorial
- Published: 2026-08-08

---

**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).

```python
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.

```python
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.

```python
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.

```python
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.

```python
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.

```python
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.

```python
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.

```python

# 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.