110 lines
4.2 KiB
Python
110 lines
4.2 KiB
Python
#!/usr/bin/env python
|
||
# coding=utf-8
|
||
'''
|
||
Author: John
|
||
Email: johnjim0816@gmail.com
|
||
Date: 2020-11-22 23:21:53
|
||
LastEditor: John
|
||
LastEditTime: 2020-11-24 19:52:40
|
||
Discription:
|
||
Environment:
|
||
'''
|
||
from itertools import count
|
||
import torch
|
||
import os
|
||
from torch.utils.tensorboard import SummaryWriter
|
||
|
||
from env import env_init
|
||
from params import get_args
|
||
from agent import PolicyGradient
|
||
from params import SEQUENCE, SAVED_MODEL_PATH, RESULT_PATH
|
||
from utils import save_results,save_model
|
||
from plot import plot
|
||
def train(cfg):
|
||
env,state_dim,n_actions = env_init()
|
||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 检测gpu
|
||
agent = PolicyGradient(state_dim,device = device,lr = cfg.policy_lr)
|
||
'''下面带pool都是存放的transition序列用于gradient'''
|
||
state_pool = [] # 存放每batch_size个episode的state序列
|
||
action_pool = []
|
||
reward_pool = []
|
||
''' 存储每个episode的reward用于绘图'''
|
||
rewards = []
|
||
moving_average_rewards = []
|
||
log_dir=os.path.split(os.path.abspath(__file__))[0]+"/logs/train/" + SEQUENCE
|
||
writer = SummaryWriter(log_dir) # 使用tensorboard的writer
|
||
for i_episode in range(cfg.train_eps):
|
||
state = env.reset()
|
||
ep_reward = 0
|
||
for _ in count():
|
||
action = agent.choose_action(state) # 根据当前环境state选择action
|
||
next_state, reward, done, _ = env.step(action)
|
||
ep_reward += reward
|
||
if done:
|
||
reward = 0
|
||
state_pool.append(state)
|
||
action_pool.append(float(action))
|
||
reward_pool.append(reward)
|
||
state = next_state
|
||
if done:
|
||
print('Episode:', i_episode, ' Reward:', ep_reward)
|
||
break
|
||
if i_episode > 0 and i_episode % cfg.batch_size == 0:
|
||
agent.update(reward_pool,state_pool,action_pool)
|
||
state_pool = [] # 每个episode的state
|
||
action_pool = []
|
||
reward_pool = []
|
||
rewards.append(ep_reward)
|
||
if i_episode == 0:
|
||
moving_average_rewards.append(ep_reward)
|
||
else:
|
||
moving_average_rewards.append(
|
||
0.9*moving_average_rewards[-1]+0.1*ep_reward)
|
||
writer.add_scalars('rewards',{'raw':rewards[-1], 'moving_average': moving_average_rewards[-1]}, i_episode+1)
|
||
writer.close()
|
||
print('Complete training！')
|
||
save_model(agent,model_path=SAVED_MODEL_PATH)
|
||
'''存储reward等相关结果'''
|
||
save_results(rewards,moving_average_rewards,tag='train',result_path=RESULT_PATH)
|
||
plot(rewards)
|
||
plot(moving_average_rewards,ylabel='moving_average_rewards_train')
|
||
|
||
def eval(cfg,saved_model_path = SAVED_MODEL_PATH):
|
||
env,state_dim,n_actions = env_init()
|
||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 检测gpu
|
||
agent = PolicyGradient(state_dim,device = device,lr = cfg.policy_lr)
|
||
agent.load_model(saved_model_path+'checkpoint.pth')
|
||
rewards = []
|
||
moving_average_rewards = []
|
||
log_dir=os.path.split(os.path.abspath(__file__))[0]+"/logs/eval/" + SEQUENCE
|
||
writer = SummaryWriter(log_dir) # 使用tensorboard的writer
|
||
for i_episode in range(cfg.eval_eps):
|
||
state = env.reset()
|
||
ep_reward = 0
|
||
for _ in count():
|
||
action = agent.choose_action(state) # 根据当前环境state选择action
|
||
next_state, reward, done, _ = env.step(action)
|
||
ep_reward += reward
|
||
state = next_state
|
||
if done:
|
||
print('Episode:', i_episode, ' Reward:', ep_reward)
|
||
break
|
||
rewards.append(ep_reward)
|
||
if i_episode == 0:
|
||
moving_average_rewards.append(ep_reward)
|
||
else:
|
||
moving_average_rewards.append(
|
||
0.9*moving_average_rewards[-1]+0.1*ep_reward)
|
||
writer.add_scalars('rewards',{'raw':rewards[-1], 'moving_average': moving_average_rewards[-1]}, i_episode+1)
|
||
writer.close()
|
||
print('Complete evaling！')
|
||
|
||
if __name__ == "__main__":
|
||
cfg = get_args()
|
||
if cfg.train:
|
||
train(cfg)
|
||
eval(cfg)
|
||
else:
|
||
model_path = os.path.split(os.path.abspath(__file__))[0]+"/saved_model/"
|
||
eval(cfg,saved_model_path=model_path)
|