[Python] DQN learning Cartpole reward graph

I've tried to make a DQN model for the simple cartpole game, but after training it for almost 3000 episodes, it produces a really weird reward graph and I'm not sure if it's even training correctly.

I've seen other posts where people didn't have a target network, didn't wait to train the model until they've had enough previous experiences and other issues. I've looked and tried to address all these issues, and this is what I got.



Here's the code:

dqn.py:

import torch
from torch import nn
import torch.nn.functional as F

class DQN(nn.Module):
def __init__(self, state_dim, action_dim, hidden_dim=256):
super(DQN, self).__init__()

self.fc1 = nn.Linear(state_dim, hidden_dim)
self.fc2 = nn.Linear(hidden_dim, action_dim)

def forward(self, x: torch.Tensor) -> torch.Tensor:
x = F.relu(self.fc1(x))
return self.fc2(x)


agent.py:

import torch
import flappy_bird_gymnasium
import gymnasium
from dqn import DQN
from experience_replay import ReplayMemory
import itertools
import yaml
import random
from pathlib import Path
from torch import nn
import matplotlib.pyplot as plt
import argparse
import numpy as np

device = "gpu" if torch.cuda.is_available() else "cpu"

class Agent():

def __init__(self, hyperparameter_set):
with open(str(Path().resolve()) + "\\hyperparameters.yml", "r") as file:
all_hyperparameters_set = yaml.safe_load(file)
hyperparameters = all_hyperparameters_set[hyperparameter_set]

self.env_name = hyperparameters["env_id"]
self.replay_memory_size = hyperparameters['replay_memory_size']
self.mini_batch_size = hyperparameters['mini_batch_size']
self.epsilon_init = hyperparameters['epsilon_init']
self.decay_rate = hyperparameters['epsilon_decay']
self.epsilon_min = hyperparameters['epsilon_min']
self.network_sync_rate = hyperparameters["network_sync_rate"]
self.lr = hyperparameters["lr"]
self.discount_factor_g = hyperparameters["discount_factor_g"]
self.hidden_layers = hyperparameters["hidden_layers"]
self.stop_on_reward = hyperparameters["stop_on_reward"]

self.loss_fn = nn.MSELoss()
self.optimizer = None

self.LOG_FILE = str(Path().resolve()) + "\\logs.txt"
self.MODEL_FILE = str(Path().resolve()) + "\\models.pt"
self.GRAPH_FILE = str(Path().resolve()) + "\\graphs.png"

def run(self, is_training=False, render=False):
env = gymnasium.make(self.env_name, render_mode="human")

obs, _ = env.reset()

reward_per_episode = []

num_actions = env.action_space.n
num_states = env.observation_space.shape[0]

policy_dqn = DQN(num_states, num_actions, self.hidden_layers)

if is_training:
policy_dqn.train()
memory = ReplayMemory(self.replay_memory_size)

epsilon_history = []
epsilon = self.epsilon_init

self.optimizer = torch.optim.Adam(policy_dqn.parameters(), lr=self.lr)

target_dqn = DQN(num_states, num_actions, self.hidden_layers).to(device)
target_dqn.load_state_dict(policy_dqn.state_dict())

step_count = 0

best_reward = float('-inf')
else:
policy_dqn.eval()
policy_dqn.load_state_dict(torch.load(self.MODEL_FILE))

for episode in itertools.count():
state, _ = env.reset()
state = torch.tensor(state, dtype=torch.float, device=device)
episode_reward = 0.0
terminated = False

while (not terminated) and episode_reward < self.stop_on_reward:

if is_training and random.random() < epsilon:
action = env.action_space.sample()
action = torch.tensor(action, dtype=torch.int64, device=device)
else:
with torch.no_grad():
action = torch.argmax(torch.squeeze(policy_dqn(torch.unsqueeze(state, 0)), 0))

new_state, reward, terminated, _, info = env.step(action.item())
new_state = torch.tensor(new_state, dtype=torch.float, device=device)

episode_reward += reward

reward = torch.tensor(reward, dtype=torch.float, device=device)

if is_training:
memory.append((state, action, new_state, reward, terminated))

step_count += 1

state = new_state

if is_training:
if episode_reward > best_reward:
log_message = f"New best reward {episode_reward: 0.1f}, {(episode_reward - best_reward):0.1f} increase from previous best"
with open(self.LOG_FILE, 'a') as file:
file.write(log_message + '\n')
torch.save(policy_dqn.state_dict(), Path(self.MODEL_FILE))
best_reward = episode_reward

if len(memory) > self.mini_batch_size:
mini_batch = memory.sample(self.mini_batch_size)

self.optimize(mini_batch, policy_dqn, target_dqn)

if step_count > self.network_sync_rate:
target_dqn.load_state_dict(policy_dqn.state_dict())
step_count = 0

epsilon = max(epsilon * self.decay_rate, self.epsilon_min)
epsilon_history.append(epsilon)
print(epsilon)
reward_per_episode.append(episode_reward)

def optimize(self, mini_batch, policy_dqn, target_dqn):
states, actions, new_states, rewards, terminations = zip(*mini_batch)

states = torch.stack(states)
actions = torch.stack(actions)
new_states = torch.stack(new_states)
rewards = torch.stack(rewards)
terminations = torch.tensor(terminations).float().to(device)

with torch.no_grad():
target_q = rewards + (1 - terminations) * self.discount_factor_g * target_dqn(new_states).max(dim=1)[0]

current_q = torch.squeeze(torch.gather(policy_dqn(states), 1, torch.unsqueeze(actions, 1)), 1)

loss = self.loss_fn(current_q, target_q)

self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()

if __name__ == "__main__":
parser = argparse.ArgumentParser(description=" Train or test model.")
parser.add_argument("hyperparameters", help='')
parser.add_argument("--train", help="Training mode", action="store_true")
args = parser.parse_args()

dql = Agent(hyperparameter_set=args.hyperparameters)

if args.train:
dql.run(is_training=True)
else:
dql.run(is_training=False, render=True)


hyperparameters.yml:

cartpole1:
env_id: CartPole-v1
replay_memory_size: 100000
mini_batch_size: 32
epsilon_init: 1
epsilon_decay: 0.9995
epsilon_min: 0.05
network_sync_rate: 10
lr: 0.001
discount_factor_g: 0.99
stop_on_reward: 100000
hidden_layers: 10
enable_double_dqn: True


I've tried playing around with the hyperparameters (even copying the parameters from stable-baselines3), but it still does the same thing.

If it is normal, this training does seem excessively slow, so what would I need to do for quicker training? Increasing network_sync rate, decreasing the epsilon decay or something else. I'm really unsure.

After experimenting some more with the hyperparameters, I was able to get something much more sensible: .

I slowed down the epsilon decay, increased the learning rate and increased the number of hidden layers.

Continue reading...