diff --git a/dqn/DQN.py b/dqn/DQN.py new file mode 100644 index 0000000..2de8072 --- /dev/null +++ b/dqn/DQN.py @@ -0,0 +1,143 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +import numpy as np +import gym +import matplotlib.pyplot as plt +import copy + +# hyper-parameters +BATCH_SIZE = 128 +LR = 0.01 +GAMMA = 0.90 +EPISILO = 0.9 +MEMORY_CAPACITY = 2000 +Q_NETWORK_ITERATION = 100 + +env = gym.make("CartPole-v0") +env = env.unwrapped +NUM_ACTIONS = env.action_space.n +NUM_STATES = env.observation_space.shape[0] +ENV_A_SHAPE = 0 if isinstance(env.action_space.sample(), int) else env.action_space.sample.shape + +class Net(nn.Module): + """docstring for Net""" + def __init__(self): + super(Net, self).__init__() + self.fc1 = nn.Linear(NUM_STATES, 50) + self.fc1.weight.data.normal_(0,0.1) + self.fc2 = nn.Linear(50,30) + self.fc2.weight.data.normal_(0,0.1) + self.out = nn.Linear(30,NUM_ACTIONS) + self.out.weight.data.normal_(0,0.1) + + def forward(self,x): + x = self.fc1(x) + x = F.relu(x) + x = self.fc2(x) + x = F.relu(x) + action_prob = self.out(x) + return action_prob + +class DQN(): + """docstring for DQN""" + def __init__(self): + super(DQN, self).__init__() + self.eval_net, self.target_net = Net(), Net() + + self.learn_step_counter = 0 + self.memory_counter = 0 + self.memory = np.zeros((MEMORY_CAPACITY, NUM_STATES * 2 + 2)) + # why the NUM_STATE*2 +2 + # When we store the memory, we put the state, action, reward and next_state in the memory + # here reward and action is a number, state is a ndarray + self.optimizer = torch.optim.Adam(self.eval_net.parameters(), lr=LR) + self.loss_func = nn.MSELoss() + + def choose_action(self, state): + state = torch.unsqueeze(torch.FloatTensor(state), 0) # get a 1D array + if np.random.randn() <= EPISILO:# greedy policy + action_value = self.eval_net.forward(state) + action = torch.max(action_value, 1)[1].data.numpy() + action = action[0] if ENV_A_SHAPE == 0 else action.reshape(ENV_A_SHAPE) + else: # random policy + action = np.random.randint(0,NUM_ACTIONS) + action = action if ENV_A_SHAPE ==0 else action.reshape(ENV_A_SHAPE) + return action + + + def store_transition(self, state, action, reward, next_state): + transition = np.hstack((state, [action, reward], next_state)) + index = self.memory_counter % MEMORY_CAPACITY + self.memory[index, :] = transition + self.memory_counter += 1 + + + def learn(self): + + #update the parameters + if self.learn_step_counter % Q_NETWORK_ITERATION ==0: + self.target_net.load_state_dict(self.eval_net.state_dict()) + self.learn_step_counter+=1 + + #sample batch from memory + sample_index = np.random.choice(MEMORY_CAPACITY, BATCH_SIZE) + batch_memory = self.memory[sample_index, :] + batch_state = torch.FloatTensor(batch_memory[:, :NUM_STATES]) + batch_action = torch.LongTensor(batch_memory[:, NUM_STATES:NUM_STATES+1].astype(int)) + batch_reward = torch.FloatTensor(batch_memory[:, NUM_STATES+1:NUM_STATES+2]) + batch_next_state = torch.FloatTensor(batch_memory[:,-NUM_STATES:]) + + #q_eval + q_eval = self.eval_net(batch_state).gather(1, batch_action) + q_next = self.target_net(batch_next_state).detach() + q_target = batch_reward + GAMMA * q_next.max(1)[0].view(BATCH_SIZE, 1) + loss = self.loss_func(q_eval, q_target) + + self.optimizer.zero_grad() + loss.backward() + self.optimizer.step() + +def reward_func(env, x, x_dot, theta, theta_dot): + r1 = (env.x_threshold - abs(x))/env.x_threshold - 0.5 + r2 = (env.theta_threshold_radians - abs(theta)) / env.theta_threshold_radians - 0.5 + reward = r1 + r2 + return reward + +def main(): + dqn = DQN() + episodes = 400 + print("Collecting Experience....") + reward_list = [] + plt.ion() + fig, ax = plt.subplots() + for i in range(episodes): + state = env.reset() + ep_reward = 0 + while True: + env.render() + action = dqn.choose_action(state) + next_state, _ , done, info = env.step(action) + x, x_dot, theta, theta_dot = next_state + reward = reward_func(env, x, x_dot, theta, theta_dot) + + dqn.store_transition(state, action, reward, next_state) + ep_reward += reward + + if dqn.memory_counter >= MEMORY_CAPACITY: + dqn.learn() + if done: + print("episode: {} , the episode reward is {}".format(i, round(ep_reward, 3))) + if done: + break + state = next_state + r = copy.copy(reward) + reward_list.append(r) + ax.set_xlim(0,300) + #ax.cla() + ax.plot(reward_list, 'g-', label='total_loss') + plt.pause(0.001) + + +if __name__ == '__main__': + main()