main.py

import numpy as np
import pandas as pd
import argparse
from utils import choose_env, args
from gym_modifly.env import DroneEnv

from torch import nn as nn
from stable_baselines3 import PPO, A2C, TD3, SAC, DDPG
from sb3_contrib import TRPO, RecurrentPPO
from stable_baselines3.common.vec_env import VecNormalize
from stable_baselines3.common.noise import NormalActionNoise, OrnsteinUhlenbeckActionNoise

if __name__ == '__main__':
    args = args()

    net_arch = {
        "small": dict(pi=[64, 64], qf=[64, 64]),
        "medium": dict(pi=[256, 256], qf=[256, 256]),
    }
    
    timesteps = 1000000
    
    if args.algo == 'randomm':
        env = DroneEnv()
        env.reset()
        for i in range(timesteps):
            action = env.action_space.sample()
            obs, reward, terminate, truncate, _ = env.step(action)
            if terminate or truncate:
                env.reset()
        
        env.close()
    else:
        env, env_test = choose_env()

        if args.algo == 'ppo':
            model = PPO("MlpPolicy", env, gamma=args.gamma, learning_rate=args.alpha,
                batch_size=args.batch_size, n_steps=args.steps, ent_coef=args.ent_coef, clip_range=args.clip_range,
                n_epochs=args.epochs, vf_coef=args.vf_coef,
                policy_kwargs=dict(net_arch=net_arch[args.arch]),
            )
            model.learn(total_timesteps=timesteps)
            env.save("vec_normalize_btc.pkl")
            # model.save("PPO_Model")

            env_test = VecNormalize.load("vec_normalize_btc.pkl", env_test)
            #  do not update them at test time
            env_test.training = False
            # reward normalization is not needed at test time
            env_test.norm_reward = False

            obs = env_test.reset()
            for _ in range(10000):
                action, _states = model.predict(obs)
                obs, rewards, dones, info = env_test.step(action)
            
            env.close()
            env_test.close()
        
        elif args.algo == 'rppo':
            model = RecurrentPPO("MlpLstmPolicy", env, gamma=args.gamma, learning_rate=args.alpha,
                batch_size=args.batch_size, n_steps=args.steps,
            )
            model.learn(total_timesteps=timesteps)
            env.save("vec_normalize_btc.pkl")
            # model.save("PPO_Model")

            env_test = VecNormalize.load("vec_normalize_btc.pkl", env_test)
            #  do not update them at test time
            env_test.training = False
            # reward normalization is not needed at test time
            env_test.norm_reward = False

            obs = env_test.reset()
            for _ in range(5000):
                action, _states = model.predict(obs)
                obs, rewards, dones, info = env_test.step(action)
            
            env.close()
            env_test.close()
        
        elif args.algo == 'ddpg':
            # The noise objects for DDPG
            n_actions = env.action_space.shape[-1]
            action_noise = NormalActionNoise(mean=np.zeros(n_actions), sigma=0.1 * np.ones(n_actions))
            
            model = DDPG("MlpPolicy", env, action_noise=action_noise, learning_rate=args.alpha, batch_size=args.batch_size, 
                         learning_starts=args.learning_starts, policy_kwargs=dict(net_arch=net_arch[args.arch]),
            )
            model.learn(total_timesteps=timesteps)
            env.save("vec_normalize_btc.pkl")

            env_test = VecNormalize.load("vec_normalize_btc.pkl", env_test)
            #  do not update them at test time
            env_test.training = False
            # reward normalization is not needed at test time
            env_test.norm_reward = False

            obs = env_test.reset()
            for _ in range(5000):
                action, _states = model.predict(obs)
                obs, rewards, dones, info = env_test.step(action)
            
            env.close()
            env_test.close()

        elif args.algo == 'sac':
            model = SAC("MlpPolicy", env, gamma=args.gamma, learning_rate=args.alpha,
                batch_size=args.batch_size, learning_starts=args.learning_starts,
                policy_kwargs=dict(net_arch=net_arch[args.arch]),
            )
            model.learn(total_timesteps=timesteps)
            env.save("vec_normalize_btc.pkl")

            env_test = VecNormalize.load("vec_normalize_btc.pkl", env_test)
            #  do not update them at test time
            env_test.training = False
            # reward normalization is not needed at test time
            env_test.norm_reward = False

            obs = env_test.reset()
            for _ in range(5000):
                action, _states = model.predict(obs)
                obs, rewards, dones, info = env_test.step(action)
            
            env.close()
            env_test.close()

        elif args.algo == 'td3':
            # The noise objects for TD3
            n_actions = env.action_space.shape[-1]
            action_noise = NormalActionNoise(mean=np.zeros(n_actions), sigma=0.1 * np.ones(n_actions))

            model = TD3("MlpPolicy", env, action_noise=action_noise, gamma=args.gamma, learning_rate=args.alpha,
                batch_size=args.batch_size, learning_starts=args.learning_starts,
                policy_kwargs=dict(net_arch=net_arch[args.arch]),
            )
            model.learn(total_timesteps=timesteps)
            env.save("vec_normalize_btc.pkl")

            env_test = VecNormalize.load("vec_normalize_btc.pkl", env_test)
            #  do not update them at test time
            env_test.training = False
            # reward normalization is not needed at test time
            env_test.norm_reward = False

            obs = env_test.reset()
            for _ in range(5000):
                action, _states = model.predict(obs)
                obs, rewards, dones, info = env_test.step(action)
            
            env.close()
            env_test.close()
            
        elif args.algo == 'a2c':
            model = A2C("MlpPolicy", env, verbose=0)
            
            model.learn(total_timesteps=timesteps)
            env.save("vec_normalize_btc.pkl")

            #testing
            env_test = VecNormalize.load("vec_normalize_btc.pkl", env_test)
            #  do not update them at test time
            env_test.training = False
            # reward normalization is not needed at test time
            env_test.norm_reward = False

            obs = env_test.reset()
            for _ in range(5000):
                action, _states = model.predict(obs)
                obs, rewards, dones, info = env_test.step(action)
            
            env.close()
            env_test.close()
        
        elif args.algo == 'trpo':
            model = TRPO("MlpPolicy", env, gamma=args.gamma, learning_rate=args.alpha,
                batch_size=args.batch_size, n_steps=args.steps,
                policy_kwargs=dict(net_arch=net_arch[args.arch]),
            )
            model.learn(total_timesteps=timesteps) 
            env.save("vec_normalize_btc.pkl")

            env_test = VecNormalize.load("vec_normalize_btc.pkl", env_test)
            #  do not update them at test time
            env_test.training = False
            # reward normalization is not needed at test time
            env_test.norm_reward = False

            obs = env_test.reset()
            for _ in range(5000):
                action, _states = model.predict(obs)
                obs, rewards, dones, info = env_test.step(action)
            
            env.close()
            env_test.close()
        
        else:
            print('Invalid Algorithm')