# Implementing Reinforcement Learning with OpenAI Gym in PyTorch: A Complete Guide

> Learn to implement reinforcement learning with OpenAI Gym in PyTorch. Build a policy network and train it using the REINFORCE algorithm for effective agent training.

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

---

**You can implement reinforcement learning in PyTorch by creating a policy network that maps environment states to action probabilities, then training it with the REINFORCE algorithm using experience collected from OpenAI Gym environments.**

The **microsoft/AI-For-Beginners** curriculum provides a production-ready implementation of deep reinforcement learning through a complete PyTorch example. This guide explains implementing reinforcement learning with OpenAI Gym in PyTorch using the repository’s `CartPole-RL-PyTorch.ipynb` notebook, which demonstrates policy gradient methods for training autonomous agents.

## Core Architecture Components

The implementation follows the standard actor-critic pattern where a neural network acts as the policy function, directly learning action probabilities from raw environment states.

### Policy Network Design

The solution uses a simple feed-forward architecture defined in `torch.nn.Module`. According to the source code in `lessons/6-Other/22-DeepRL/CartPole-RL-PyTorch.ipynb`, the network consists of two fully-connected layers with ReLU activation. The forward pass accepts observation vectors and returns raw logits for each possible action, which are then converted to probabilities using `torch.softmax`.

### Environment Interface

While the notebook imports **gymnasium** (the maintained successor to OpenAI Gym), the API remains fully compatible with standard Gym environments. The code initializes environments using `gym.make()` and processes the standard five-tuple return `(state, reward, terminated, truncated, info)` from each `env.step()` call.

## The REINFORCE Training Loop

The policy gradient method updates network weights based on the likelihood of actions that resulted in high cumulative rewards.

### Experience Collection

During each episode, the agent samples actions from a categorical distribution created via `torch.distributions.Categorical(probs)`. The code stores `log_prob` tensors and received rewards in separate lists until the episode terminates or truncates. This separation allows the algorithm to compute returns retrospectively after the complete trajectory is observed.

### Policy Gradient Update

The algorithm computes discounted returns using the formula $G_t = r_t + \gamma G_{t+1}$, then applies the REINFORCE loss: `-(log_prob * return).mean()`. This loss is backpropagated through the network and optimized using `torch.optim.Adam` with a configurable learning rate (referenced as `alpha` in the curriculum).

## Complete Implementation Example

Below is a minimal, self-contained script that mirrors the repository’s approach. This implementation works with any Gym-compatible environment by adjusting the `obs_dim` and `act_dim` variables.

```python
import gymnasium as gym
import torch
import torch.nn as nn
import torch.optim as optim

# ---------- 1. Define the policy network ----------

class PolicyNet(nn.Module):
    def __init__(self, obs_dim, act_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(obs_dim, 128),
            nn.ReLU(),
            nn.Linear(128, act_dim),
        )

    def forward(self, x):
        return self.net(x)

# ---------- 2. Hyper-parameters ----------

env_name = "CartPole-v1"          # replace with any Gym env

num_episodes = 500
gamma = 0.99                      # discount factor

learning_rate = 1e-3

# ---------- 3. Initialise environment & policy ----------

env = gym.make(env_name)
obs_dim = env.observation_space.shape[0]
act_dim = env.action_space.n

policy = PolicyNet(obs_dim, act_dim)
optimizer = optim.Adam(policy.parameters(), lr=learning_rate)

# ---------- 4. REINFORCE training loop ----------

for episode in range(num_episodes):
    state, _ = env.reset()
    log_probs = []
    rewards = []

    done = False
    while not done:
        state_tensor = torch.from_numpy(state).float()
        logits = policy(state_tensor)
        probs = torch.softmax(logits, dim=-1)
        action_dist = torch.distributions.Categorical(probs)
        action = action_dist.sample()
        log_probs.append(action_dist.log_prob(action))

        state, reward, terminated, truncated, _ = env.step(action.item())
        done = terminated or truncated
        rewards.append(reward)

    # Compute discounted returns

    returns = []
    G = 0.0
    for r in reversed(rewards):
        G = r + gamma * G
        returns.insert(0, G)
    returns = torch.tensor(returns)

    # Policy gradient update

    loss = -(torch.stack(log_probs) * returns).mean()
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if (episode + 1) % 50 == 0:
        total_reward = sum(rewards)
        print(f"Episode {episode+1:3d} – Return: {total_reward:.1f}")

print("Training complete!")

```

## Adapting to Custom Gym Environments

You can apply this same pattern to any Gym-compatible environment by following these modifications:

- **Install dependencies**: Ensure `gymnasium` (or `gym`) and `torch` are listed in your [`environment.yml`](https://github.com/microsoft/AI-For-Beginners/blob/main/environment.yml) or [`requirements.txt`](https://github.com/microsoft/AI-For-Beginners/blob/main/requirements.txt) as shown in the repository root files.

- **Adjust dimensions**: Modify `obs_dim` to match `env.observation_space.shape[0]` and `act_dim` to match `env.action_space.n` for discrete environments.

- **Handle continuous spaces**: For continuous action spaces, replace the `Categorical` distribution with `torch.distributions.Normal` and adjust the network output to predict mean and standard deviation parameters.

## Key Files in the Repository

Understanding the codebase structure helps when extending the implementation:

- **`lessons/6-Other/22-DeepRL/CartPole-RL-PyTorch.ipynb`**: The full Jupyter notebook containing the complete REINFORCE implementation with visualization code.

- **[`environment.yml`](https://github.com/microsoft/AI-For-Beginners/blob/main/environment.yml)**: Conda environment definition specifying Python 3.11, PyTorch, and `gymnasium` dependencies.

- **[`requirements.txt`](https://github.com/microsoft/AI-For-Beginners/blob/main/requirements.txt)**: Root-level dependency list for the entire curriculum.

- **[`examples/README.md`](https://github.com/microsoft/AI-For-Beginners/blob/main/examples/README.md)**: Quick-start guide for running Python scripts within the repository framework.

## Summary

- The **microsoft/AI-For-Beginners** repository provides a complete REINFORCE implementation in `CartPole-RL-PyTorch.ipynb` using PyTorch and Gymnasium.

- The architecture requires only a simple two-layer network, a categorical distribution for sampling, and the policy gradient loss `-(log_prob * return).mean()`.

- Training involves collecting full episodes before computing discounted returns and performing optimizer steps with `torch.optim.Adam`.

- The same code structure works for any Gym environment by adjusting observation dimensions and action space configurations.

## Frequently Asked Questions

### Does this implementation work with the original OpenAI Gym or only Gymnasium?

The code works with both. The repository uses **gymnasium** (the maintained fork of OpenAI Gym), but the API is identical. You can replace `import gymnasium as gym` with `import gym` if using the original package, though Gymnasium is recommended for bug fixes and continued support.

### What hyperparameters should I tune for more complex environments?

Start by adjusting the **learning rate** (typically between `1e-4` and `1e-2` for Adam) and the **discount factor** `gamma` (usually `0.95` to `0.99`). For high-dimensional state spaces, increase the hidden layer size beyond 128 units or add additional layers, ensuring you modify the `PolicyNet` class architecture accordingly.

### How do I know if my policy is converging?

Monitor the **episode return** (sum of rewards) as shown in the logging code. In the CartPole environment, convergence is typically indicated when the return reaches 500 (the maximum episode length) consistently across multiple consecutive episodes. For other environments, look for a stable upward trend in the rolling average return over 100 episodes.

### Can I use this code for continuous action spaces like MuJoCo environments?

No, not without modification. The provided `PolicyNet` outputs discrete action probabilities using `Categorical`. For continuous actions, you must change the network to output mean and log-standard-deviation parameters, then use `torch.distributions.Normal` for sampling. The REINFORCE loss formula remains the same, but you will backpropagate through the Gaussian log probabilities instead.