hot update A2C

projects/codes/A2C/README.md

@@ -1,5 +0,0 @@
-## A2C
-
-
-
-https://towardsdatascience.com/understanding-actor-critic-methods-931b97b6df3f

projects/codes/A2C/a2c.py

@@ -1,56 +1,60 @@
-#!/usr/bin/env python
-# coding=utf-8
-'''
-Author: JiangJi
-Email: johnjim0816@gmail.com
-Date: 2021-05-03 22:16:08
-LastEditor: JiangJi
-LastEditTime: 2022-07-20 23:54:40
-Discription: 
-Environment: 
-'''
 import torch
-import torch.optim as optim
-import torch.nn as nn
-import torch.nn.functional as F
-from torch.distributions import Categorical
+import numpy as np
+
 
-class ActorCritic(nn.Module):
-    ''' A2C网络模型，包含一个Actor和Critic
-    '''
-    def __init__(self, input_dim, output_dim, hidden_dim):
-        super(ActorCritic, self).__init__()
-        self.critic = nn.Sequential(
-            nn.Linear(input_dim, hidden_dim),
-            nn.ReLU(),
-            nn.Linear(hidden_dim, 1)
-        )
         
-        self.actor = nn.Sequential(
-            nn.Linear(input_dim, hidden_dim),
-            nn.ReLU(),
-            nn.Linear(hidden_dim, output_dim),
-            nn.Softmax(dim=1),
-        )
-        
-    def forward(self, x):
-        value = self.critic(x)
-        probs = self.actor(x)
-        dist  = Categorical(probs)
-        return dist, value
 class A2C:
-    ''' A2C算法
-    '''
-    def __init__(self,n_states,n_actions,cfg) -> None:
-        self.gamma = cfg.gamma
-        self.device = torch.device(cfg.device)
-        self.model = ActorCritic(n_states, n_actions, cfg.hidden_size).to(self.device)
-        self.optimizer = optim.Adam(self.model.parameters())
+    def __init__(self,models,memories,cfg):
+        self.n_actions = cfg['n_actions']
+        self.gamma = cfg['gamma']
+        self.device = torch.device(cfg['device']) 
+        self.memory = memories['ACMemory']
+        self.actor = models['Actor'].to(self.device)
+        self.critic = models['Critic'].to(self.device)
+        self.actor_optim = torch.optim.Adam(self.actor.parameters(), lr=cfg['actor_lr'])
+        self.critic_optim = torch.optim.Adam(self.critic.parameters(), lr=cfg['critic_lr'])
+    def sample_action(self,state):
+        state = torch.tensor(state, device=self.device, dtype=torch.float32).unsqueeze(dim=0)
+        dist = self.actor(state)
+        value = self.critic(state) # note that 'dist' need require_grad=True
+        value = value.detach().numpy().squeeze(0)[0]
+        action = np.random.choice(self.n_actions, p=dist.detach().numpy().squeeze(0)) # shape(p=(n_actions,1)
+        return action,value,dist 
+    def predict_action(self,state):
+        state = torch.tensor(state, device=self.device, dtype=torch.float32).unsqueeze(dim=0)
+        dist = self.actor(state)
+        value = self.critic(state) # note that 'dist' need require_grad=True
+        value = value.detach().numpy().squeeze(0)[0]
+        action = np.random.choice(self.n_actions, p=dist.detach().numpy().squeeze(0)) # shape(p=(n_actions,1)
+        return action,value,dist 
+    def update(self,next_state,entropy):
+        value_pool,log_prob_pool,reward_pool = self.memory.sample()
+        next_state = torch.tensor(next_state, device=self.device, dtype=torch.float32).unsqueeze(dim=0)
+        next_value = self.critic(next_state)
+        returns = np.zeros_like(reward_pool)
+        for t in reversed(range(len(reward_pool))):
+            next_value = reward_pool[t] + self.gamma * next_value # G(s_{t},a{t}) = r_{t+1} + gamma * V(s_{t+1})
+            returns[t] = next_value
+        returns = torch.tensor(returns, device=self.device)
+        value_pool = torch.tensor(value_pool, device=self.device)
+        advantages = returns - value_pool
+        log_prob_pool = torch.stack(log_prob_pool)
+        actor_loss = (-log_prob_pool * advantages).mean()
+        critic_loss = 0.5 * advantages.pow(2).mean()
+        tot_loss = actor_loss + critic_loss + 0.001 * entropy
+        self.actor_optim.zero_grad()
+        self.critic_optim.zero_grad()
+        tot_loss.backward()
+        self.actor_optim.step()
+        self.critic_optim.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.actor.state_dict(), f"{path}/actor_checkpoint.pt")
+        torch.save(self.critic.state_dict(), f"{path}/critic_checkpoint.pt")
 
-    def compute_returns(self,next_value, rewards, masks):
-        R = next_value
-        returns = []
-        for step in reversed(range(len(rewards))):
-            R = rewards[step] + self.gamma * R * masks[step]
-            returns.insert(0, R)
-        return returns
+    def load_model(self, path):
+        self.actor.load_state_dict(torch.load(f"{path}/actor_checkpoint.pt"))
+        self.critic.load_state_dict(torch.load(f"{path}/critic_checkpoint.pt"))

projects/codes/A2C/a2c_2.py

@@ -0,0 +1,55 @@
+import torch
+import numpy as np
+
+class A2C_2:
+    def __init__(self,models,memories,cfg):
+        self.n_actions = cfg['n_actions']
+        self.gamma = cfg['gamma']
+        self.device = torch.device(cfg['device']) 
+        self.memory = memories['ACMemory']
+        self.ac_net = models['ActorCritic'].to(self.device)
+        self.ac_optimizer = torch.optim.Adam(self.ac_net.parameters(), lr=cfg['lr'])
+    def sample_action(self,state):
+        state = torch.tensor(state, device=self.device, dtype=torch.float32).unsqueeze(dim=0)
+        value, dist = self.ac_net(state) # note that 'dist' need require_grad=True
+        value = value.detach().numpy().squeeze(0)[0]
+        action = np.random.choice(self.n_actions, p=dist.detach().numpy().squeeze(0)) # shape(p=(n_actions,1)
+        return action,value,dist
+    def predict_action(self,state):
+        ''' predict can be all wrapped with no_grad(), then donot need detach(), or you can just copy contents of 'sample_action'
+        '''
+        with torch.no_grad(): 
+            state = torch.tensor(state, device=self.device, dtype=torch.float32).unsqueeze(dim=0)
+            value, dist = self.ac_net(state)
+            value = value.numpy().squeeze(0)[0] # shape(value) = (1,)
+            action = np.random.choice(self.n_actions, p=dist.numpy().squeeze(0)) # shape(p=(n_actions,1)
+        return action,value,dist
+    def update(self,next_state,entropy):
+        value_pool,log_prob_pool,reward_pool = self.memory.sample()
+        next_state = torch.tensor(next_state, device=self.device, dtype=torch.float32).unsqueeze(dim=0)
+        next_value,_ = self.ac_net(next_state)
+        returns = np.zeros_like(reward_pool)
+        for t in reversed(range(len(reward_pool))):
+            next_value = reward_pool[t] + self.gamma * next_value # G(s_{t},a{t}) = r_{t+1} + gamma * V(s_{t+1})
+            returns[t] = next_value
+        returns = torch.tensor(returns, device=self.device)
+        value_pool = torch.tensor(value_pool, device=self.device)
+        advantages = returns - value_pool
+        log_prob_pool = torch.stack(log_prob_pool)
+        actor_loss = (-log_prob_pool * advantages).mean()
+        critic_loss = 0.5 * advantages.pow(2).mean()
+        ac_loss = actor_loss + critic_loss + 0.001 * entropy
+        self.ac_optimizer.zero_grad()
+        ac_loss.backward()
+        self.ac_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.ac_net.state_dict(), f"{path}/a2c_checkpoint.pt")
+
+    def load_model(self, path):
+        self.ac_net.load_state_dict(torch.load(f"{path}/a2c_checkpoint.pt"))
+        
+       

projects/codes/A2C/main.py

@@ -0,0 +1,121 @@
+import sys,os
+os.environ["KMP_DUPLICATE_LIB_OK"]  =  "TRUE" # avoid "OMP: Error #15: Initializing libiomp5md.dll, but found libiomp5md.dll already initialized."
+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 datetime
+import argparse
+import gym
+import torch
+import numpy as np
+from common.utils import all_seed
+from common.launcher import Launcher
+from common.memories import PGReplay
+from common.models import ActorSoftmax,Critic
+from envs.register import register_env
+from a2c import A2C
+
+class Main(Launcher):
+    def get_args(self):
+        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='A2C',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=1600,type=int,help="episodes of training") 
+        parser.add_argument('--test_eps',default=20,type=int,help="episodes of testing") 
+        parser.add_argument('--ep_max_steps',default = 100000,type=int,help="steps per episode, much larger value can simulate infinite steps")
+        parser.add_argument('--gamma',default=0.99,type=float,help="discounted factor") 
+        parser.add_argument('--actor_lr',default=3e-4,type=float,help="learning rate of actor")
+        parser.add_argument('--critic_lr',default=1e-3,type=float,help="learning rate of critic")
+        parser.add_argument('--actor_hidden_dim',default=256,type=int,help="hidden of actor net")
+        parser.add_argument('--critic_hidden_dim',default=256,type=int,help="hidden of critic net")
+        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('--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()   
+        default_args = {'result_path':f"{curr_path}/outputs/{args.env_name}/{curr_time}/results/",
+                        'model_path':f"{curr_path}/outputs/{args.env_name}/{curr_time}/models/",
+        }
+        args = {**vars(args),**default_args}  # type(dict)                         
+        return args
+    def env_agent_config(self,cfg):
+        ''' create env and agent
+        '''  
+        register_env(cfg['env_name'])
+        env = gym.make(cfg['env_name']) 
+        if cfg['seed'] !=0: # set random seed
+            all_seed(env,seed=cfg["seed"]) 
+        try: # state dimension
+            n_states = env.observation_space.n # print(hasattr(env.observation_space, 'n'))
+        except AttributeError:
+            n_states = env.observation_space.shape[0] # print(hasattr(env.observation_space, 'shape'))
+        n_actions = env.action_space.n  # action dimension
+        print(f"n_states: {n_states}, n_actions: {n_actions}")
+        cfg.update({"n_states":n_states,"n_actions":n_actions}) # update to cfg paramters
+        models = {'Actor':ActorSoftmax(cfg['n_states'],cfg['n_actions'], hidden_dim = cfg['actor_hidden_dim']),'Critic':Critic(cfg['n_states'],1,hidden_dim=cfg['critic_hidden_dim'])}
+        memories = {'ACMemory':PGReplay()}
+        agent = A2C(models,memories,cfg)
+        return env,agent
+    def train(self,cfg,env,agent):
+        print("Start training!")
+        print(f"Env: {cfg['env_name']}, Algorithm: {cfg['algo_name']}, Device: {cfg['device']}")
+        rewards = []  # record rewards for all episodes
+        steps = [] # record steps for all episodes
+        
+        for i_ep in range(cfg['train_eps']):
+            ep_reward = 0  # reward per episode
+            ep_step = 0 # step per episode
+            ep_entropy = 0
+            state = env.reset()  # reset and obtain initial state
+            
+            for _ in range(cfg['ep_max_steps']):
+                action, value, dist = agent.sample_action(state)  # sample action
+                next_state, reward, done, _ = env.step(action)  # update env and return transitions
+                log_prob = torch.log(dist.squeeze(0)[action])
+                entropy = -np.sum(np.mean(dist.detach().numpy()) * np.log(dist.detach().numpy()))
+                agent.memory.push((value,log_prob,reward))  # save transitions
+                state = next_state  # update state
+                ep_reward += reward
+                ep_entropy += entropy
+                ep_step += 1
+                if done:
+                    break
+            agent.update(next_state,ep_entropy)  # update agent
+            rewards.append(ep_reward)
+            steps.append(ep_step)
+            if (i_ep+1)%10==0:
+                print(f'Episode: {i_ep+1}/{cfg["train_eps"]}, Reward: {ep_reward:.2f}, Steps:{ep_step}')
+        print("Finish training!")
+        return {'episodes':range(len(rewards)),'rewards':rewards,'steps':steps}
+    def test(self,cfg,env,agent):
+        print("Start testing!")
+        print(f"Env: {cfg['env_name']}, Algorithm: {cfg['algo_name']}, Device: {cfg['device']}")
+        rewards = []  # record rewards for all episodes
+        steps = [] # record steps for all episodes
+        for i_ep in range(cfg['test_eps']):
+            ep_reward = 0  # reward per episode
+            ep_step = 0
+            state = env.reset()  # reset and obtain initial state
+            for _ in range(cfg['ep_max_steps']):
+                action,_,_ = agent.predict_action(state)  # predict action
+                next_state, reward, done, _ = env.step(action)  
+                state = next_state 
+                ep_reward += reward
+                ep_step += 1
+                if done:
+                    break
+            rewards.append(ep_reward)
+            steps.append(ep_step)
+            print(f"Episode: {i_ep+1}/{cfg['test_eps']}, Steps:{ep_step}, Reward: {ep_reward:.2f}")
+        print("Finish testing!")
+        return {'episodes':range(len(rewards)),'rewards':rewards,'steps':steps}
+
+if __name__ == "__main__":
+    main = Main()
+    main.run()
+   
+
+        
+    

projects/codes/A2C/main2.py

@@ -0,0 +1,120 @@
+import sys,os
+os.environ["KMP_DUPLICATE_LIB_OK"]  =  "TRUE" # avoid "OMP: Error #15: Initializing libiomp5md.dll, but found libiomp5md.dll already initialized."
+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 datetime
+import argparse
+import gym
+import torch
+import numpy as np
+from common.utils import all_seed
+from common.launcher import Launcher
+from common.memories import PGReplay
+from common.models import ActorCriticSoftmax
+from envs.register import register_env
+from a2c_2 import A2C_2
+
+class Main(Launcher):
+    def get_args(self):
+        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='A2C',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=2000,type=int,help="episodes of training") 
+        parser.add_argument('--test_eps',default=20,type=int,help="episodes of testing") 
+        parser.add_argument('--ep_max_steps',default = 100000,type=int,help="steps per episode, much larger value can simulate infinite steps")
+        parser.add_argument('--gamma',default=0.99,type=float,help="discounted factor") 
+        parser.add_argument('--lr',default=3e-4,type=float,help="learning rate")
+        parser.add_argument('--actor_hidden_dim',default=256,type=int)
+        parser.add_argument('--critic_hidden_dim',default=256,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('--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()   
+        default_args = {'result_path':f"{curr_path}/outputs/{args.env_name}/{curr_time}/results/",
+                        'model_path':f"{curr_path}/outputs/{args.env_name}/{curr_time}/models/",
+        }
+        args = {**vars(args),**default_args}  # type(dict)                         
+        return args
+    def env_agent_config(self,cfg):
+        ''' create env and agent
+        '''  
+        register_env(cfg['env_name'])
+        env = gym.make(cfg['env_name']) 
+        if cfg['seed'] !=0: # set random seed
+            all_seed(env,seed=cfg["seed"]) 
+        try: # state dimension
+            n_states = env.observation_space.n # print(hasattr(env.observation_space, 'n'))
+        except AttributeError:
+            n_states = env.observation_space.shape[0] # print(hasattr(env.observation_space, 'shape'))
+        n_actions = env.action_space.n  # action dimension
+        print(f"n_states: {n_states}, n_actions: {n_actions}")
+        cfg.update({"n_states":n_states,"n_actions":n_actions}) # update to cfg paramters
+        models = {'ActorCritic':ActorCriticSoftmax(cfg['n_states'],cfg['n_actions'], actor_hidden_dim = cfg['actor_hidden_dim'],critic_hidden_dim=cfg['critic_hidden_dim'])}
+        memories = {'ACMemory':PGReplay()}
+        agent = A2C_2(models,memories,cfg)
+        return env,agent
+    def train(self,cfg,env,agent):
+        print("Start training!")
+        print(f"Env: {cfg['env_name']}, Algorithm: {cfg['algo_name']}, Device: {cfg['device']}")
+        rewards = []  # record rewards for all episodes
+        steps = [] # record steps for all episodes
+        
+        for i_ep in range(cfg['train_eps']):
+            ep_reward = 0  # reward per episode
+            ep_step = 0 # step per episode
+            ep_entropy = 0
+            state = env.reset()  # reset and obtain initial state
+            
+            for _ in range(cfg['ep_max_steps']):
+                action, value, dist = agent.sample_action(state)  # sample action
+                next_state, reward, done, _ = env.step(action)  # update env and return transitions
+                log_prob = torch.log(dist.squeeze(0)[action])
+                entropy = -np.sum(np.mean(dist.detach().numpy()) * np.log(dist.detach().numpy()))
+                agent.memory.push((value,log_prob,reward))  # save transitions
+                state = next_state  # update state
+                ep_reward += reward
+                ep_entropy += entropy
+                ep_step += 1
+                if done:
+                    break
+            agent.update(next_state,ep_entropy)  # update agent
+            rewards.append(ep_reward)
+            steps.append(ep_step)
+            if (i_ep+1)%10==0:
+                print(f'Episode: {i_ep+1}/{cfg["train_eps"]}, Reward: {ep_reward:.2f}, Steps:{ep_step}')
+        print("Finish training!")
+        return {'episodes':range(len(rewards)),'rewards':rewards,'steps':steps}
+    def test(self,cfg,env,agent):
+        print("Start testing!")
+        print(f"Env: {cfg['env_name']}, Algorithm: {cfg['algo_name']}, Device: {cfg['device']}")
+        rewards = []  # record rewards for all episodes
+        steps = [] # record steps for all episodes
+        for i_ep in range(cfg['test_eps']):
+            ep_reward = 0  # reward per episode
+            ep_step = 0
+            state = env.reset()  # reset and obtain initial state
+            for _ in range(cfg['ep_max_steps']):
+                action,_,_ = agent.predict_action(state)  # predict action
+                next_state, reward, done, _ = env.step(action)  
+                state = next_state 
+                ep_reward += reward
+                ep_step += 1
+                if done:
+                    break
+            rewards.append(ep_reward)
+            steps.append(ep_step)
+            print(f"Episode: {i_ep+1}/{cfg['test_eps']}, Steps:{ep_step}, Reward: {ep_reward:.2f}")
+        print("Finish testing!")
+        return {'episodes':range(len(rewards)),'rewards':rewards,'steps':steps}
+
+if __name__ == "__main__":
+    main = Main()
+    main.run()
+   
+
+        
+    

projects/codes/A2C/outputs/CartPole-v0/20220713-221850/results/params.json

@@ -1,14 +0,0 @@
-{
-    "algo_name": "A2C",
-    "env_name": "CartPole-v0",
-    "n_envs": 8,
-    "max_steps": 20000,
-    "n_steps": 5,
-    "gamma": 0.99,
-    "lr": 0.001,
-    "hidden_dim": 256,
-    "deivce": "cpu",
-    "result_path": "C:\\Users\\24438\\Desktop\\rl-tutorials/outputs/CartPole-v0/20220713-221850/results/",
-    "model_path": "C:\\Users\\24438\\Desktop\\rl-tutorials/outputs/CartPole-v0/20220713-221850/models/",
-    "save_fig": true
-}

projects/codes/A2C/outputs/CartPole-v0/20220829-135818/results/params.json

@@ -0,0 +1,19 @@
+{
+    "algo_name": "A2C",
+    "env_name": "CartPole-v0",
+    "train_eps": 2000,
+    "test_eps": 20,
+    "ep_max_steps": 100000,
+    "gamma": 0.99,
+    "lr": 0.0003,
+    "actor_hidden_dim": 256,
+    "critic_hidden_dim": 256,
+    "device": "cpu",
+    "seed": 10,
+    "show_fig": false,
+    "save_fig": true,
+    "result_path": "/Users/jj/Desktop/rl-tutorials/codes/A2C/outputs/CartPole-v0/20220829-135818/results/",
+    "model_path": "/Users/jj/Desktop/rl-tutorials/codes/A2C/outputs/CartPole-v0/20220829-135818/models/",
+    "n_states": 4,
+    "n_actions": 2
+}

projects/codes/A2C/outputs/CartPole-v0/20220829-135818/results/testing_results.csv

@@ -0,0 +1,21 @@
+episodes,rewards,steps
+0,200.0,200
+1,200.0,200
+2,93.0,93
+3,155.0,155
+4,116.0,116
+5,200.0,200
+6,190.0,190
+7,176.0,176
+8,200.0,200
+9,200.0,200
+10,200.0,200
+11,179.0,179
+12,200.0,200
+13,185.0,185
+14,191.0,191
+15,200.0,200
+16,200.0,200
+17,124.0,124
+18,200.0,200
+19,172.0,172

projects/codes/A2C/outputs/CartPole-v0/20220829-143327/results/params.json

@@ -0,0 +1 @@
+{"algo_name": "A2C", "env_name": "CartPole-v0", "train_eps": 1600, "test_eps": 20, "ep_max_steps": 100000, "gamma": 0.99, "actor_lr": 0.0003, "critic_lr": 0.001, "actor_hidden_dim": 256, "critic_hidden_dim": 256, "device": "cpu", "seed": 10, "show_fig": false, "save_fig": true, "result_path": "/Users/jj/Desktop/rl-tutorials/codes/A2C/outputs/CartPole-v0/20220829-143327/results/", "model_path": "/Users/jj/Desktop/rl-tutorials/codes/A2C/outputs/CartPole-v0/20220829-143327/models/", "n_states": 4, "n_actions": 2}

projects/codes/A2C/outputs/CartPole-v0/20220829-143327/results/testing_results.csv

@@ -0,0 +1,21 @@
+episodes,rewards,steps
+0,177.0,177
+1,180.0,180
+2,200.0,200
+3,200.0,200
+4,167.0,167
+5,124.0,124
+6,128.0,128
+7,200.0,200
+8,200.0,200
+9,200.0,200
+10,186.0,186
+11,187.0,187
+12,200.0,200
+13,176.0,176
+14,200.0,200
+15,200.0,200
+16,200.0,200
+17,200.0,200
+18,185.0,185
+19,180.0,180

projects/codes/A2C/task0.py

@@ -1,137 +0,0 @@
-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 to system path
-
-import gym
-import numpy as np
-import torch
-import torch.optim as optim
-import datetime
-import argparse
-from common.multiprocessing_env import SubprocVecEnv
-from a2c import ActorCritic
-from common.utils import save_results, make_dir
-from common.utils import plot_rewards, save_args
-
-
-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='A2C',type=str,help="name of algorithm")
-    parser.add_argument('--env_name',default='CartPole-v0',type=str,help="name of environment")
-    parser.add_argument('--n_envs',default=8,type=int,help="numbers of environments")
-
-    parser.add_argument('--max_steps',default=20000,type=int,help="episodes of training")
-    parser.add_argument('--n_steps',default=5,type=int,help="episodes of testing")
-    parser.add_argument('--gamma',default=0.99,type=float,help="discounted factor")
-    parser.add_argument('--lr',default=1e-3,type=float,help="learning rate")
-    parser.add_argument('--hidden_dim',default=256,type=int)
-    parser.add_argument('--device',default='cpu',type=str,help="cpu or cuda") 
-    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/' ) # path to save models
-    parser.add_argument('--save_fig',default=True,type=bool,help="if save figure or not")           
-    args = parser.parse_args()                        
-    return args
-
-def make_envs(env_name):
-    def _thunk():
-        env = gym.make(env_name)
-        env.seed(2)
-        return env
-    return _thunk
-def test_env(env,model,vis=False):
-    state = env.reset()
-    if vis: env.render()
-    done = False
-    total_reward = 0
-    while not done:
-        state = torch.FloatTensor(state).unsqueeze(0).to(cfg.device)
-        dist, _ = model(state)
-        next_state, reward, done, _ = env.step(dist.sample().cpu().numpy()[0])
-        state = next_state
-        if vis: env.render()
-        total_reward += reward
-    return total_reward
-
-def compute_returns(next_value, rewards, masks, gamma=0.99):
-    R = next_value
-    returns = []
-    for step in reversed(range(len(rewards))):
-        R = rewards[step] + gamma * R * masks[step]
-        returns.insert(0, R)
-    return returns
-
-
-def train(cfg,envs):
-    print('Start training!')
-    print(f'Env:{cfg.env_name}, Algorithm:{cfg.algo_name}, Device:{cfg.device}')
-    env = gym.make(cfg.env_name) # a single env
-    env.seed(10)
-    n_states  = envs.observation_space.shape[0]
-    n_actions = envs.action_space.n
-    model = ActorCritic(n_states, n_actions, cfg.hidden_dim).to(cfg.device)
-    optimizer = optim.Adam(model.parameters())
-    step_idx    = 0
-    test_rewards = []
-    test_ma_rewards = []
-    state = envs.reset()
-    while step_idx < cfg.max_steps:
-        log_probs = []
-        values    = []
-        rewards   = []
-        masks     = []
-        entropy = 0
-        # rollout trajectory
-        for _ in range(cfg.n_steps):
-            state = torch.FloatTensor(state).to(cfg.device)
-            dist, value = model(state)
-            action = dist.sample()
-            next_state, reward, done, _ = envs.step(action.cpu().numpy())
-            log_prob = dist.log_prob(action)
-            entropy += dist.entropy().mean()
-            log_probs.append(log_prob)
-            values.append(value)
-            rewards.append(torch.FloatTensor(reward).unsqueeze(1).to(cfg.device))
-            masks.append(torch.FloatTensor(1 - done).unsqueeze(1).to(cfg.device))
-            state = next_state
-            step_idx += 1
-            if step_idx % 100 == 0:
-                test_reward = np.mean([test_env(env,model) for _ in range(10)])
-                print(f"step_idx:{step_idx}, test_reward:{test_reward}")
-                test_rewards.append(test_reward)
-                if test_ma_rewards:
-                    test_ma_rewards.append(0.9*test_ma_rewards[-1]+0.1*test_reward)
-                else:
-                    test_ma_rewards.append(test_reward) 
-                # plot(step_idx, test_rewards)   
-        next_state = torch.FloatTensor(next_state).to(cfg.device)
-        _, next_value = model(next_state)
-        returns = compute_returns(next_value, rewards, masks)
-        log_probs = torch.cat(log_probs)
-        returns   = torch.cat(returns).detach()
-        values    = torch.cat(values)
-        advantage = returns - values
-        actor_loss  = -(log_probs * advantage.detach()).mean()
-        critic_loss = advantage.pow(2).mean()
-        loss = actor_loss + 0.5 * critic_loss - 0.001 * entropy
-        optimizer.zero_grad()
-        loss.backward()
-        optimizer.step()
-    print('Finish training！')
-    return {'rewards':test_rewards,'ma_rewards':test_ma_rewards}
-if __name__ == "__main__":
-    cfg = get_args()
-    envs = [make_envs(cfg.env_name) for i in range(cfg.n_envs)]
-    envs = SubprocVecEnv(envs) 
-    # training
-    res_dic = train(cfg,envs)
-    make_dir(cfg.result_path,cfg.model_path)
-    save_args(cfg)
-    save_results(res_dic, tag='train',
-                 path=cfg.result_path)
-    plot_rewards(res_dic['rewards'], res_dic['ma_rewards'], cfg, tag="train") # 画出结果