From 4126da321a6cbe44a64884f649671fecf0fd3c46 Mon Sep 17 00:00:00 2001 From: Johnny He <269401927@qq.com> Date: Thu, 18 Oct 2018 16:51:16 +0800 Subject: [PATCH] Create a2c_cartpole.py --- A2C/a2c_cartpole.py | 155 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 155 insertions(+) create mode 100644 A2C/a2c_cartpole.py diff --git a/A2C/a2c_cartpole.py b/A2C/a2c_cartpole.py new file mode 100644 index 0000000..a415a1b --- /dev/null +++ b/A2C/a2c_cartpole.py @@ -0,0 +1,155 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.autograd import Variable +import matplotlib.pyplot as plt +import numpy as np +import math +import random +import os +import gym +import warnings +warnings.filterwarnings('ignore') + +# Hyper Parameters +STATE_DIM = 4 +ACTION_DIM = 2 +STEP = 2000 +SAMPLE_NUMS = 30 + + +class ActorNetwork(nn.Module): + + def __init__(self,input_size,hidden_size,action_size): + super(ActorNetwork, self).__init__() + self.fc1 = nn.Linear(input_size,hidden_size) + self.fc2 = nn.Linear(hidden_size,hidden_size) + self.fc3 = nn.Linear(hidden_size,action_size) + + def forward(self,x): + out = F.relu(self.fc1(x)) + out = F.relu(self.fc2(out)) + out = F.log_softmax(self.fc3(out), dim=-1) + return out + +class ValueNetwork(nn.Module): + + def __init__(self,input_size,hidden_size,output_size): + super(ValueNetwork, self).__init__() + self.fc1 = nn.Linear(input_size,hidden_size) + self.fc2 = nn.Linear(hidden_size,hidden_size) + self.fc3 = nn.Linear(hidden_size,output_size) + + def forward(self,x): + out = F.relu(self.fc1(x)) + out = F.relu(self.fc2(out)) + out = self.fc3(out) + return out + +def roll_out(actor_network,task,sample_nums,value_network,init_state): + #task.reset() + states = [] + actions = [] + rewards = [] + is_done = False + final_r = 0 + state = init_state + + for j in range(sample_nums): + states.append(state) + log_softmax_action = actor_network(Variable(torch.Tensor([state]))) + softmax_action = torch.exp(log_softmax_action) + action = np.random.choice(ACTION_DIM,p=softmax_action.cpu().data.numpy()[0]) + one_hot_action = [int(k == action) for k in range(ACTION_DIM)] + next_state,reward,done,_ = task.step(action) + #fix_reward = -10 if done else 1 + actions.append(one_hot_action) + rewards.append(reward) + final_state = next_state + state = next_state + if done: + is_done = True + state = task.reset() + break + if not is_done: + final_r = value_network(Variable(torch.Tensor([final_state]))).cpu().data.numpy() + + return states,actions,rewards,final_r,state + +def discount_reward(r, gamma,final_r): + discounted_r = np.zeros_like(r) + running_add = final_r + for t in reversed(range(0, len(r))): + running_add = running_add * gamma + r[t] + discounted_r[t] = running_add + return discounted_r + +def main(): + # init a task generator for data fetching + task = gym.make("CartPole-v0") + init_state = task.reset() + + # init value network + value_network = ValueNetwork(input_size = STATE_DIM,hidden_size = 40,output_size = 1) + value_network_optim = torch.optim.Adam(value_network.parameters(),lr=0.01) + + # init actor network + actor_network = ActorNetwork(STATE_DIM,40,ACTION_DIM) + actor_network_optim = torch.optim.Adam(actor_network.parameters(),lr = 0.01) + + steps =[] + task_episodes =[] + test_results =[] + + for step in range(STEP): + states,actions,rewards,final_r,current_state = roll_out(actor_network,task,SAMPLE_NUMS,value_network,init_state) + init_state = current_state + actions_var = Variable(torch.Tensor(actions).view(-1,ACTION_DIM)) + states_var = Variable(torch.Tensor(states).view(-1,STATE_DIM)) + + # train actor network + actor_network_optim.zero_grad() + log_softmax_actions = actor_network(states_var) + vs = value_network(states_var).detach() + # calculate qs + qs = Variable(torch.Tensor(discount_reward(rewards,0.99,final_r))) + + advantages = qs - vs + actor_network_loss = - torch.mean(torch.sum(log_softmax_actions*actions_var,1)* advantages) + actor_network_loss.backward() + torch.nn.utils.clip_grad_norm(actor_network.parameters(),0.5) + actor_network_optim.step() + + # train value network + value_network_optim.zero_grad() + target_values = qs + values = value_network(states_var) + criterion = nn.MSELoss() + + target_values = target_values.view(target_values.shape[0], 1) + value_network_loss = criterion(values, target_values) + value_network_loss.backward() + torch.nn.utils.clip_grad_norm(value_network.parameters(),0.5) + value_network_optim.step() + + # Testing + if (step + 1) % 50== 0: + result = 0 + test_task = gym.make("CartPole-v0") + for test_epi in range(10): + state = test_task.reset() + for test_step in range(200): + softmax_action = torch.exp(actor_network(Variable(torch.Tensor([state])))) + #print(softmax_action.data) + action = np.argmax(softmax_action.data.numpy()[0]) + next_state,reward,done,_ = test_task.step(action) + result += reward + state = next_state + if done: + break + print("step:",step+1,"test result:",result/10.0) + steps.append(step+1) + test_results.append(result/10) + +if __name__ == '__main__': + main()