ASSISTANT
Sure! Below is a step-by-step guide to creating a simple Deep Q-Network (DQN) agent to solve the **CartPole-v1** environment from OpenAI's Gym library. The agent will learn to balance the pole, and after training, you can visualize its performance and plot the rewards over time.
### **Prerequisites**
Ensure you have the following Python packages installed:
- `gym` for the environment
- `numpy` for numerical operations
- `matplotlib` for plotting
- `torch` (PyTorch) for building the neural network
You can install them using `pip`:
```bash
pip install gym numpy matplotlib torch
```
### **Complete Code**
```python
import gym
import numpy as np
import random
from collections import deque
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim
# Set seeds for reproducibility
seed = 42
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
# Define the neural network model
class DQN(nn.Module):
def __init__(self, state_size, action_size, hidden_size=24):
super(DQN, self).__init__()
self.fc1 = nn.Linear(state_size, hidden_size)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(hidden_size, hidden_size)
self.fc3 = nn.Linear(hidden_size, action_size)
def forward(self, x):
out = self.relu(self.fc1(x))
out = self.relu(self.fc2(out))
return self.fc3(out)
# Replay Buffer to store experience tuples
class ReplayBuffer:
def __init__(self, capacity):
self.buffer = deque(maxlen=capacity)
def push(self, state, action, reward, next_state, done):
self.buffer.append( (state, action, reward, next_state, done) )
def sample(self, batch_size):
return random.sample(self.buffer, batch_size)
def __len__(self):
return len(self.buffer)
# DQN Agent
class DQNAgent:
def __init__(self, state_size, action_size, hidden_size=24,
lr=1e-3, gamma=0.99, epsilon=1.0, epsilon_min=0.01,
epsilon_decay=0.995, batch_size=64, buffer_capacity=10000,
target_update=10):
self.state_size = state_size
self.action_size = action_size
self.hidden_size = hidden_size
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Networks
self.policy_net = DQN(state_size, action_size, hidden_size).to(self.device)
self.target_net = DQN(state_size, action_size, hidden_size).to(self.device)
self.target_net.load_state_dict(self.policy_net.state_dict())
self.target_net.eval() # Target network is not trained
self.optimizer = optim.Adam(self.policy_net.parameters(), lr=lr)
self.loss_fn = nn.MSELoss()
# Replay buffer
self.memory = ReplayBuffer(buffer_capacity)
self.batch_size = batch_size
# Hyperparameters
self.gamma = gamma
self.epsilon = epsilon
self.epsilon_min = epsilon_min
self.epsilon_decay = epsilon_decay
self.target_update = target_update
self.steps_done = 0
def select_action(self, state):
self.steps_done += 1
if random.random() < self.epsilon:
return random.randrange(self.action_size)
else:
state = torch.FloatTensor(state).unsqueeze(0).to(self.device)
with torch.no_grad():
q_values = self.policy_net(state)
return q_values.argmax().item()
def push_memory(self, state, action, reward, next_state, done):
self.memory.push(state, action, reward, next_state, done)
def learn(self):
if len(self.memory) < self.batch_size:
return
batch = self.memory.sample(self.batch_size)
states, actions, rewards, next_states, dones = zip(*batch)
states = torch.FloatTensor(states).to(self.device)
actions = torch.LongTensor(actions).unsqueeze(1).to(self.device)
rewards = torch.FloatTensor(rewards).unsqueeze(1).to(self.device)
next_states = torch.FloatTensor(next_states).to(self.device)
dones = torch.FloatTensor(dones).unsqueeze(1).to(self.device)
# Current Q values
q_values = self.policy_net(states).gather(1, actions)
# Next Q values from target network
with torch.no_grad():
next_q_values = self.target_net(next_states).max(1)[0].unsqueeze(1)
target_q = rewards + (1 - dones) * self.gamma * next_q_values
# Compute loss
loss = self.loss_fn(q_values, target_q)
# Backpropagation
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
# Update epsilon
if self.epsilon > self.epsilon_min:
self.epsilon *= self.epsilon_decay
def update_target_network(self):
self.target_net.load_state_dict(self.policy_net.state_dict())
# Training the agent
def train_dqn(env, agent, episodes=500):
rewards = []
for episode in range(1, episodes+1):
state = env.reset()
done = False
total_reward = 0
while not done:
action = agent.select_action(state)
next_state, reward, done, _ = env.step(action)
agent.push_memory(state, action, reward, next_state, done)
agent.learn()
state = next_state
total_reward += reward
rewards.append(total_reward)
# Update target network
if episode % agent.target_update == 0:
agent.update_target_network()
# Print progress
if episode % 10 == 0:
avg_reward = np.mean(rewards[-10:])
print(f"Episode {episode}, Average Reward: {avg_reward:.2f}, Epsilon: {agent.epsilon:.2f}")
return rewards
# Evaluation: Let the agent play and render
def evaluate_agent(env, agent, episodes=5):
agent.epsilon = 0.0 # No exploration
for episode in range(1, episodes+1):
state = env.reset()
done = False
total_reward = 0
while not done:
env.render()
action = agent.select_action(state)
state, reward, done, _ = env.step(action)
total_reward += reward
print(f"Evaluation Episode {episode}: Total Reward: {total_reward}")
env.close()
# Plotting the rewards
def plot_rewards(rewards):
plt.figure(figsize=(12,5))
plt.plot(rewards, label='Reward per Episode')
# Moving average for smoother curve
ma = np.convolve(rewards, np.ones((10,))/10, mode='valid')
plt.plot(ma, label='Moving Average (10 episodes)')
plt.xlabel('Episode')
plt.ylabel('Total Reward')
plt.title('DQN Agent Performance on CartPole-v1')
plt.legend()
plt.show()
# Main function to run everything
def main():
env = gym.make('CartPole-v1')
env.seed(seed)
state_size = env.observation_space.shape[0]
action_size = env.action_space.n
agent = DQNAgent(state_size, action_size)
print("Training started...")
rewards = train_dqn(env, agent, episodes=500)
print("Training finished!")
plot_rewards(rewards)
print("Evaluating the agent...")
evaluate_agent(env, agent, episodes=5)
if __name__ == "__main__":
main()
```
### **Explanation of the Code**
1. **Imports and Seeding:**
- Import necessary libraries.
- Set seeds for reproducibility.
2. **Neural Network (`DQN`):**
- A simple feedforward neural network with two hidden layers using ReLU activation.
- Maps the state to Q-values for each action.
3. **Replay Buffer (`ReplayBuffer`):**
- Stores experience tuples `(state, action, reward, next_state, done)`.
- Allows the agent to learn from past experiences by sampling random batches.
4. **DQN Agent (`DQNAgent`):**
- Initializes policy and target networks.
- Implements the ε-greedy policy for exploration vs. exploitation.
- Handles the learning process by sampling from the replay buffer and updating the network weights.
- Periodically updates the target network to stabilize training.
5. **Training Function (`train_dqn`):**
- Runs the training loop for a specified number of episodes.
- Collects rewards for each episode for later visualization.
- Prints the average reward every 10 episodes to monitor progress.
6. **Evaluation Function (`evaluate_agent`):**
- Runs the agent without exploration (ε=0) to demonstrate its learned policy.
- Renders the environment so you can watch the agent in action.
7. **Plotting Function (`plot_rewards`):**
- Plots the total reward per episode.
- Includes a moving average to visualize performance trends more clearly.
8. **Main Function (`main`):**
- Sets up the environment and agent.
- Triggers the training, plotting, and evaluation processes.
### **Running the Code**
1. **Save the Code:**
Save the above code in a Python file, e.g., `dqn_cartpole.py`.
2. **Execute the Script:**
Run the script using Python:
```bash
python dqn_cartpole.py
```
3. **Observe Training Progress:**
The script will print the average reward every 10 episodes. As training progresses, you should see the average reward increasing, indicating that the agent is learning to balance the pole better.
4. **View the Plot:**
After training, a plot will display the rewards over episodes, showing how the agent's performance improves over time.
5. **Watch the Agent:**
The evaluation phase will render the environment, allowing you to watch the trained agent perform. You can see how it balances the pole without exploration.
### **Sample Output**
```
Training started...
Episode 10, Average Reward: 10.30, Epsilon: 0.60
Episode 20, Average Reward: 20.50, Epsilon: 0.36
...
Episode 500, Average Reward: 500.00, Epsilon: 0.01
Training finished!
```
A plot will appear showing the rewards per episode, typically increasing over time as the agent learns. Then, the rendered environment window will display the agent successfully balancing the pole.
### **Notes**
- **Hyperparameters:** The chosen hyperparameters (like learning rate, discount factor, etc.) work well for CartPole but might need tuning for other environments.
- **Performance:** Due to the simplicity of the environment and the agent, training should be relatively fast. However, for more complex environments, more sophisticated techniques and longer training times may be necessary.
- **Extensions:** You can experiment by modifying the network architecture, using different exploration strategies, implementing Prioritized Experience Replay, or using other RL algorithms like Double DQN or Dueling DQN for improved performance.
Happy Reinforcement Learning!