How to Build Deep Reinforcement Learning Agents with PyTorch: A Complete CartPole Tutorial

You can build deep reinforcement learning agents with PyTorch by implementing a policy gradient (REINFORCE) algorithm that uses a neural network to map environment states to action probabilities, training the agent to maximize cumulative rewards through interaction with the OpenAI Gym environment.

Deep reinforcement learning combines neural networks with the classic RL paradigm to let an agent learn optimal behavior through interaction with an environment. In the Microsoft AI-For-Beginners curriculum, the CartPole-RL-PyTorch.ipynb notebook demonstrates a complete end-to-end workflow for creating a deep RL agent using PyTorch and the OpenAI Gym simulator. This guide walks through the core architecture, the REINFORCE algorithm implementation, and the training loop required to build your first agent.

Core Architecture Components

Building a deep RL agent requires three essential components: the environment simulator, the policy network, and the learning algorithm.

The CartPole Environment

The environment setup uses gym.make("CartPole-v1") from OpenAI Gym, which provides a classic control problem where a pole must be balanced on a moving cart. According to the source code in lessons/6-Other/22-DeepRL/CartPole-RL-PyTorch.ipynb (lines 44-50), this environment supplies:

• An action space with 2 discrete actions (push left or right)
• An observation space with 4-dimensional state (cart position, cart velocity, pole angle, pole angular velocity)
• A reward of +1 for every timestep the pole remains balanced

Policy Network Architecture

The policy network is implemented as a simple feed-forward neural network using torch.nn.Sequential. As defined in the notebook (lines 13-18), the architecture maps the 4-dimensional state to a probability distribution over the two actions:

import torch.nn as nn

policy = nn.Sequential(
    nn.Linear(4, 128),    # Input layer: state dim → hidden dim

    nn.ReLU(),            # Non-linearity

    nn.Linear(128, 2),    # Hidden → action space

    nn.Softmax(dim=1)     # Output probability distribution

)

This network learns to generalize from raw observations to action probabilities, allowing the agent to select actions in states it has never explicitly encountered during training.

The REINFORCE Algorithm Implementation

The agent follows the REINFORCE algorithm, a policy gradient method that directly optimizes the stochastic policy without requiring a value-function estimator.

Discounted Rewards Calculation

A critical component is the discounted returns calculation. As implemented in CartPole-RL-PyTorch.ipynb (lines 85-93), the function computes the cumulative reward for each timestep, discounting future rewards by a factor γ ≈ 0.99 to emphasize long-term success:

import numpy as np

def discounted_rewards(rewards, gamma=0.99):
    ret, cum = [], 0.0
    for r in reversed(rewards):
        cum = r + gamma * cum
        ret.insert(0, cum)
    ret = np.array(ret)
    # Normalize returns for training stability

    return (ret - ret.mean()) / (ret.std() + 1e-8)

By weighting earlier rewards less, the algorithm prioritizes actions that lead to longer episode durations, which is essential for balancing tasks where immediate rewards are uniform.

Policy Gradient Loss

The loss function combines the log probabilities of selected actions with their corresponding returns. For each episode, the loss equals the negative sum of log π(a|s) multiplied by (return – baseline), which is minimized via gradient ascent on the expected return using the Adam optimizer.

Implementing the Training Loop

The training routine in the notebook (lines 100-130) follows a standard episodic structure across approximately 300 episodes:

1. Episode Collection: Run a complete episode, collecting states, actions, and rewards until termination
2. Return Computation: Calculate discounted returns for the trajectory
3. Loss Calculation: Compute policy gradient loss using action log-probabilities and normalized returns
4. Optimization Step: Perform optimizer.zero_grad(), loss.backward(), and optimizer.step()

The training loop demonstrates how the policy improves over time, with episode rewards typically increasing from near-random performance (~20-50 steps) to successfully balancing the pole for 500 steps (the environment maximum).

Complete Working Example

Below is a minimal, self-contained implementation adapted from the notebook that captures the complete deep RL agent:

import gym
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim

# Initialize environment

env = gym.make("CartPole-v1")
num_inputs = env.observation_space.shape[0]  # 4

num_actions = env.action_space.n             # 2

# Define policy network

policy = nn.Sequential(
    nn.Linear(num_inputs, 128), nn.ReLU(),
    nn.Linear(128, num_actions), nn.Softmax(dim=1)
)

optimizer = optim.Adam(policy.parameters(), lr=1e-3)

def discounted_rewards(rewards, gamma=0.99):
    ret, cum = [], 0.0
    for r in reversed(rewards):
        cum = r + gamma * cum
        ret.insert(0, cum)
    ret = np.array(ret)
    return (ret - ret.mean()) / (ret.std() + 1e-8)

def run_episode(render=False):
    state = env.reset()
    states, actions, rewards = [], [], []
    done = False
    while not done:
        if render:
            env.render()
        s = torch.from_numpy(state).float().unsqueeze(0)
        probs = policy(s)
        action = np.random.choice(num_actions, 
                                  p=probs.squeeze().detach().numpy())
        next_state, reward, done, _ = env.step(action)
        states.append(state)
        actions.append(action)
        rewards.append(reward)
        state = next_state
    return np.array(states), np.array(actions), np.array(rewards)

# Training loop

for episode in range(300):
    states, actions, rewards = run_episode()
    returns = torch.tensor(discounted_rewards(rewards), dtype=torch.float32)
    logits = policy(torch.from_numpy(states).float())
    log_probs = torch.log(logits.gather(1, torch.tensor(actions).unsqueeze(1)))
    loss = -(log_probs.squeeze() * returns).mean()
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    if episode % 50 == 0:
        print(f"Episode {episode}, total reward {rewards.sum():.1f}")

This implementation demonstrates the core mechanics of how to build deep reinforcement learning agents with PyTorch, from environment interaction to gradient-based policy updates.

Extending the Approach

Once you master the basic CartPole implementation, you can extend this foundation to more sophisticated algorithms and environments.

Actor-Critic Methods

For improved sample efficiency, combine a value (critic) network with the policy (actor). The repository's CartPole-RL-TF.ipynb (TensorFlow version) includes an actor-critic example that can be translated to PyTorch, where the critic estimates state values to reduce variance in the policy gradient.

Complex Environments

To scale beyond CartPole, replace the environment with more challenging Gym tasks such as LunarLander or Atari. This typically requires:

• Increasing network depth (adding convolutional layers for image inputs)
• Implementing experience replay buffers
• Adding target networks for stability (as in DQN)
• Using proximal policy optimization (PPO) for continuous action spaces

Summary

• Deep RL agents use neural networks to approximate policies that map environment states to action probabilities, implemented in lessons/6-Other/22-DeepRL/CartPole-RL-PyTorch.ipynb
• The REINFORCE algorithm updates policy parameters using gradient ascent on expected returns, requiring only the policy network and no explicit value function
• Discounted rewards (γ ≈ 0.99) stabilize training by weighting long-term outcomes over immediate rewards
• The CartPole-v1 environment provides a canonical testbed with 4-dimensional observations and 2 discrete actions
• Extensions to Actor-Critic methods and complex environments follow the same PyTorch patterns demonstrated in the base implementation

Frequently Asked Questions

What is the REINFORCE algorithm?

REINFORCE is a Monte Carlo policy gradient algorithm that directly optimizes the parameters of a stochastic policy to maximize expected cumulative reward. It works by running complete episodes, calculating the total return for each action taken, and updating the policy to increase the probability of actions that led to higher returns.

Why use a neural network for the policy function?

A neural network serves as a universal function approximator that can generalize across states. Unlike tabular methods that store separate probabilities for each state, the network in CartPole-RL-PyTorch.ipynb (4→128→2 architecture) learns continuous representations, allowing the agent to balance the pole in slightly different positions using similar strategies.

How long does training take for a CartPole agent?

Training typically completes within 300 episodes, which translates to roughly 2-5 minutes on a standard CPU. The agent usually reaches the maximum reward of 500 steps per episode (the environment cap) well before the 300th episode, depending on the learning rate (1e-3) and random seed.

Can this PyTorch implementation handle continuous action spaces?

The current implementation handles discrete actions (2 choices) using a softmax output layer. For continuous action spaces (like robotic control), you would modify the network to output mean and standard deviation parameters for a Gaussian distribution, then sample actions from that distribution while maintaining the policy gradient framework.

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