task0.py

import sys,os
curr_path = os.path.dirname(os.path.abspath(__file__))  # current path
parent_path = os.path.dirname(curr_path)  # parent path 
sys.path.append(parent_path)  # add path to system path

import argparse
import datetime
import gym
import torch
import random
import numpy as np
import torch.nn as nn
from common.memories import ReplayBufferQue
from common.models import MLP
from utils.utils import save_results,all_seed,plot_rewards,save_args
from softq import SoftQ

def get_args():
    """ hyperparameters
    """
    curr_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")  # obtain current time
    parser = argparse.ArgumentParser(description="hyperparameters")      
    parser.add_argument('--algo_name',default='SoftQ',type=str,help="name of algorithm")
    parser.add_argument('--env_name',default='CartPole-v0',type=str,help="name of environment")
    parser.add_argument('--train_eps',default=200,type=int,help="episodes of training")
    parser.add_argument('--test_eps',default=20,type=int,help="episodes of testing")
    parser.add_argument('--max_steps',default=200,type=int,help="maximum steps per episode")
    parser.add_argument('--gamma',default=0.99,type=float,help="discounted factor")
    parser.add_argument('--alpha',default=4,type=float,help="alpha")
    parser.add_argument('--lr',default=0.0001,type=float,help="learning rate")
    parser.add_argument('--memory_capacity',default=50000,type=int,help="memory capacity")
    parser.add_argument('--batch_size',default=128,type=int)
    parser.add_argument('--target_update',default=2,type=int)
    parser.add_argument('--device',default='cpu',type=str,help="cpu or cuda") 
    parser.add_argument('--seed',default=10,type=int,help="seed") 
    parser.add_argument('--result_path',default=curr_path + "/outputs/" + parser.parse_args().env_name + \
            '/' + curr_time + '/results/' )
    parser.add_argument('--model_path',default=curr_path + "/outputs/" + parser.parse_args().env_name + \
            '/' + curr_time + '/models/' ) 
    parser.add_argument('--show_fig',default=False,type=bool,help="if show figure or not")  
    parser.add_argument('--save_fig',default=True,type=bool,help="if save figure or not")           
    args = parser.parse_args()                          
    return args

class SoftQNetwork(nn.Module):
    '''Actually almost same to common.models.MLP
    '''
    def __init__(self,input_dim,output_dim):
        super(SoftQNetwork,self).__init__()
        self.fc1 = nn.Linear(input_dim, 64)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(64, 256)
        self.fc3 = nn.Linear(256, output_dim)
        
    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        x = self.fc3(x)
        return x

def env_agent_config(cfg):
    ''' create env and agent
    '''
    env = gym.make(cfg.env_name)  # create env
    if cfg.seed !=0: # set random seed
        all_seed(env,seed=cfg.seed)
    n_states = env.observation_space.shape[0]  # state dimension
    n_actions = env.action_space.n  # action dimension
    print(f"state dim: {n_states}, action dim: {n_actions}")
    # model = MLP(n_states,n_actions)
    model = SoftQNetwork(n_states,n_actions)
    memory =  ReplayBufferQue(cfg.memory_capacity) # replay buffer
    agent = SoftQ(n_actions,model,memory,cfg)  # create agent
    return env, agent

def train(cfg, env, agent):
    ''' training
    '''
    print("start training!")
    print(f"Env: {cfg.env_name}, Algo: {cfg.algo_name}, Device: {cfg.device}")
    rewards = []  # record rewards for all episodes
    steps = [] # record steps for all episodes, sometimes need
    for i_ep in range(cfg.train_eps):
        ep_reward = 0  # reward per episode
        ep_step = 0
        state = env.reset()   # reset and obtain initial state
        while True:
        # for _ in range(cfg.max_steps):
            ep_step += 1
            action = agent.sample_action(state) # sample action
            next_state, reward, done, _ = env.step(action)  # update env and return transitions
            agent.memory.push((state, action, reward, next_state, done))  # save transitions
            state = next_state  # update next state for env
            agent.update()  # update agent
            ep_reward += reward  
            if done:
                break
        if (i_ep + 1) % cfg.target_update == 0:  # target net update, target_update means "C" in pseucodes
            agent.target_net.load_state_dict(agent.policy_net.state_dict())
        steps.append(ep_step)
        rewards.append(ep_reward)
        if (i_ep + 1) % 10 == 0:
            print(f'Episode: {i_ep+1}/{cfg.train_eps}, Reward: {ep_reward:.2f}')
    print("finish training!")
    res_dic = {'episodes':range(len(rewards)),'rewards':rewards}
    return res_dic
def test(cfg, env, agent):
    print("start testing!")
    print(f"Env: {cfg.env_name}, Algo: {cfg.algo_name}, Device: {cfg.device}")
    rewards = []  # record rewards for all episodes
    for i_ep in range(cfg.test_eps):
        ep_reward = 0  # reward per episode
        state = env.reset()  # reset and obtain initial state
        while True:
            action = agent.predict_action(state)  # predict action
            next_state, reward, done, _ = env.step(action)  
            state = next_state  
            ep_reward += reward 
            if done:
                break
        rewards.append(ep_reward)
        print(f'Episode: {i_ep+1}/{cfg.test_eps}，Reward: {ep_reward:.2f}')
    print("finish testing!")
    env.close()
    return {'episodes':range(len(rewards)),'rewards':rewards}

if __name__ == "__main__":
    cfg = get_args()
    # 训练
    env, agent = env_agent_config(cfg)
    res_dic = train(cfg, env, agent)
    save_args(cfg,path = cfg.result_path) # 保存参数到模型路径上
    agent.save_model(path = cfg.model_path)  # 保存模型
    save_results(res_dic, tag = 'train', path = cfg.result_path)  
    plot_rewards(res_dic['rewards'], cfg, path = cfg.result_path,tag = "train")  
    # 测试
    env, agent = env_agent_config(cfg) # 也可以不加，加这一行的是为了避免训练之后环境可能会出现问题，因此新建一个环境用于测试
    agent.load_model(path = cfg.model_path)  # 导入模型
    res_dic = test(cfg, env, agent)
    save_results(res_dic, tag='test',
                 path = cfg.result_path)  # 保存结果
    plot_rewards(res_dic['rewards'], cfg, path = cfg.result_path,tag = "test")  # 画出结果