I am attempting to implement PER (Prioritized Experience Replay) along with a DQN algorithm that I've written from scratch. I've been able to solve OpenAI's cartPole-v0 environment with my DQN algorithm without PER. After implementing PER, the algorithm no longer converges. I suspect that the issue comes in when I try to anneal the bias introduced by PER with importance sampling weights. The algoritm I am trying to implement can be found on page 5 of this paper.
This is the function that is supposed to update the neural network at every iteration of the simulation:
def learn(self):
self.beta = max(self.beta, self.step_count/self.epochs)
if self.beta >= 1:
self.beta = 1
# Sample random transitions
samples = self.memory.sample(self.batch_size, self.beta)
indexes = np.array(samples['indexes'])
weights = samples['weights']
states = self.memory.data['state'][indexes]
actions = self.memory.data['action'][indexes]
rewards = self.memory.data['reward'][indexes]
terminal = self.memory.data['done'][indexes]
next_states = self.memory.data['next_state'][indexes]
states = torch.tensor(
states, dtype=torch.float32).to(self.Q_net.device)
actions = torch.tensor(
actions, dtype=torch.int64).unsqueeze(-1).to(self.Q_net.device)
rewards = torch.tensor(
rewards, dtype=torch.float32).unsqueeze(-1).to(self.Q_net.device)
terminal = torch.tensor(
terminal, dtype=torch.float32).unsqueeze(-1).to(self.Q_net.device)
next_states = torch.tensor(
next_states, dtype=torch.float32).to(self.Q_net.device)
targets = self.T_net.forward(next_states)
targets = rewards[0] + self.gamma * \
torch.max(targets, dim=1)[0] * (1-terminal)[0]
targets = targets.unsqueeze(1)
Q_s = self.Q_net.forward(states)
Q_s = torch.gather(input=Q_s, dim=1, index=actions)
priorities = (targets - Q_s).detach().numpy()
for i in range(len(priorities)):
priorities[i] = np.abs(priorities[i]) + self.min_sampling_prob
self.memory.update_priorities(indexes, priorities)
weights = torch.tensor(
weights, dtype=torch.float32).to(self.Q_net.device)
weights = torch.reshape(weights, targets.shape)
loss = nn.functional.smooth_l1_loss(weights*Q_s, targets)
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
And here is the testing script:
from environment import FullEnvironment
from dqn_agent_per import Agent
import numpy as np
import matplotlib.pyplot as plt
import gym
env = gym.make('CartPole-v0')
EPOCHS = 800
agent = Agent(env, alpha=0.5, beta=0.5, eps_start=1.0, hidden_nodes=512, batch_size=64, epochs=EPOCHS, gamma=0.999, eps_decay=0.996,
eps_min=0.1, lr=(5e-4), custom_env=False, activation='r')
eps_hist = []
avg_score = []
scores = []
TARGET_UPDATE_FREQ = 100
total_steps = 0
for i in range(EPOCHS):
state = env.reset()
episode_reward = 0.0
steps = 0
while True:
total_steps += 1
action = agent.choose_action(state)
state_next, reward, done, _ = env.step(action)
agent.store_transition(state, action, reward, state_next, done)
agent.learn()
episode_reward += reward
state = state_next
if total_steps % TARGET_UPDATE_FREQ == 0:
agent.reset_target_network()
if done:
break
eps_hist.append(agent.epsilon)
avg_score.append(episode_reward)
scores.append(np.mean(avg_score[-100:]))
agent.decay_eps()
print('Episode', i, 'score %.2f' % episode_reward, 'Avg score %.2f' %
np.mean(avg_score[-100:]), 'Epsilon %.3f' % agent.epsilon)
plt.plot(scores)
plt.show()
Thank you very much in advance!