90 lines
2.9 KiB
Python
90 lines
2.9 KiB
Python
#!/usr/bin/env python
|
|
# coding=utf-8
|
|
'''
|
|
Author: John
|
|
Email: johnjim0816@gmail.com
|
|
Date: 2020-11-22 23:27:44
|
|
LastEditor: John
|
|
LastEditTime: 2022-08-27 13:45:26
|
|
Discription:
|
|
Environment:
|
|
'''
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from torch.distributions import Bernoulli
|
|
from torch.autograd import Variable
|
|
import numpy as np
|
|
|
|
|
|
class PolicyGradient:
|
|
|
|
def __init__(self, model,memory,cfg):
|
|
self.gamma = cfg['gamma']
|
|
self.device = torch.device(cfg['device'])
|
|
self.memory = memory
|
|
self.policy_net = model.to(self.device)
|
|
self.optimizer = torch.optim.RMSprop(self.policy_net.parameters(), lr=cfg['lr'])
|
|
|
|
def sample_action(self,state):
|
|
|
|
state = torch.from_numpy(state).float()
|
|
state = Variable(state)
|
|
probs = self.policy_net(state)
|
|
print("probs")
|
|
print(probs)
|
|
m = Bernoulli(probs) # 伯努利分布
|
|
action = m.sample()
|
|
|
|
action = action.data.numpy().astype(int)[0] # 转为标量
|
|
return action
|
|
def predict_action(self,state):
|
|
|
|
state = torch.from_numpy(state).float()
|
|
state = Variable(state)
|
|
probs = self.policy_net(state)
|
|
m = Bernoulli(probs) # 伯努利分布
|
|
action = m.sample()
|
|
action = action.data.numpy().astype(int)[0] # 转为标量
|
|
return action
|
|
|
|
def update(self):
|
|
state_pool,action_pool,reward_pool= self.memory.sample()
|
|
state_pool,action_pool,reward_pool = list(state_pool),list(action_pool),list(reward_pool)
|
|
# Discount reward
|
|
running_add = 0
|
|
for i in reversed(range(len(reward_pool))):
|
|
if reward_pool[i] == 0:
|
|
running_add = 0
|
|
else:
|
|
running_add = running_add * self.gamma + reward_pool[i]
|
|
reward_pool[i] = running_add
|
|
|
|
# Normalize reward
|
|
reward_mean = np.mean(reward_pool)
|
|
reward_std = np.std(reward_pool)
|
|
for i in range(len(reward_pool)):
|
|
reward_pool[i] = (reward_pool[i] - reward_mean) / reward_std
|
|
|
|
# Gradient Desent
|
|
self.optimizer.zero_grad()
|
|
|
|
for i in range(len(reward_pool)):
|
|
state = state_pool[i]
|
|
action = Variable(torch.FloatTensor([action_pool[i]]))
|
|
reward = reward_pool[i]
|
|
state = Variable(torch.from_numpy(state).float())
|
|
probs = self.policy_net(state)
|
|
m = Bernoulli(probs)
|
|
loss = -m.log_prob(action) * reward # Negtive score function x reward
|
|
# print(loss)
|
|
loss.backward()
|
|
self.optimizer.step()
|
|
self.memory.clear()
|
|
def save_model(self,path):
|
|
from pathlib import Path
|
|
# create path
|
|
Path(path).mkdir(parents=True, exist_ok=True)
|
|
torch.save(self.policy_net.state_dict(), path+'checkpoint.pt')
|
|
def load_model(self,path):
|
|
self.policy_net.load_state_dict(torch.load(path+'checkpoint.pt')) |