<jupyter_start><jupyter_text><jupyter_code>%%capture
!pip install git+https://github.com/AI4Finance-LLC/FinRL-Library.git
!pip install wandb
# %%capture
# !pip install torch==1.4.0
import pandas as pd
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
# matplotlib.use('Agg')
import datetime

%matplotlib inline
from finrl import config
from finrl import config_tickers
from finrl.finrl_meta.preprocessor.yahoodownloader import YahooDownloader
from finrl.finrl_meta.preprocessor.preprocessors import FeatureEngineer, data_split
from finrl.finrl_meta.env_stock_trading.env_stocktrading import StockTradingEnv
from finrl.finrl_meta.env_stock_trading.env_stocktrading_np import StockTradingEnv as StockTradingEnv_numpy 
# from finrl.drl_agents.stablebaselines3.models import DRLAgent as DRLAgent_sb3
from finrl.finrl_meta.data_processor import DataProcessor

from finrl.plot import backtest_stats, backtest_plot, get_daily_return, get_baseline
import ray
from pprint import pprint
import pprint
import sys
sys.path.append("../FinRL-Library")

import itertools
import wandb
from wandb.integration.sb3 import WandbCallback
wandb.login()
import os
if not os.path.exists("./" + config.DATA_SAVE_DIR):
    os.makedirs("./" + config.DATA_SAVE_DIR)
if not os.path.exists("./" + config.TRAINED_MODEL_DIR):
    os.makedirs("./" + config.TRAINED_MODEL_DIR)
if not os.path.exists("./" + config.TENSORBOARD_LOG_DIR):
    os.makedirs("./" + config.TENSORBOARD_LOG_DIR)
if not os.path.exists("./" + config.RESULTS_DIR):
    os.makedirs("./" + config.RESULTS_DIR)
def model_params(model_name):
  sweep_config = {
      'method': 'bayes'
          }

  metric = {
      'name': 'Val sharpe',
      'goal': 'maximize'
  }

  sweep_config['metric'] = metric

  ddpg_param_dict = {
    "buffer_size": {
        "values":[int(1e4), int(1e5), int(1e6)]
        },     
    "learning_rate": {   
        "distribution": "log_uniform",
        "min": 1e-5,
        "max": 1,
    },
    "batch_size" :{
        'values':[32, 64, 128, 256, 512]
    },
  }

  a2c_param_dict = {
      "n_steps": {
          'values': [128, 256, 512, 1024, 2048]},
      "ent_coef": {   
        "distribution": "log_uniform",
        "min": 1e-8,
        "max": 1,
    },
      "learning_rate": {   
        "distribution": "log_uniform",
        "min": 1e-5,
        "max": 1,
    },
  }

  ppo_param_dict = {
      "ent_coef": {   
        "distribution": "log_uniform",
        "min": 1e-8,
        "max": 1,
    },
        "n_steps": {
            'values':[128, 256, 512, 1024, 2048]},
        "learning_rate": {   
        "distribution": "log_uniform",
        "min": 1e-2,
        "max": 1,
    },
        "batch_size": {
        'values':[32, 64, 128, 256, 512]
    },
  }

  stopping_criteria = {'type': 'hyperband', 's': 2, 'eta': 2, 'max_iter': 12}

  sweep_config['early_terminate'] = stopping_criteria

  if model_name == 'ddpg':
    sweep_config['parameters'] = ddpg_param_dict
  elif model_name == 'ppo':
    sweep_config['parameters'] = ppo_param_dict
  elif model_name == 'a2c':
    sweep_config['parameters'] = a2c_param_dict

  return sweep_config
%%writefile model_wandb.py
import wandb
from wandb.integration.sb3 import WandbCallback
import time

import numpy as np
import pandas as pd
from finrl import config
# from finrl.finrl_meta.env_stock_trading.env_stocktrading import StockTradingEnv
from finrl.finrl_meta.preprocessor.preprocessors import data_split
from stable_baselines3 import A2C, DDPG, PPO, SAC, TD3
from stable_baselines3.common.callbacks import BaseCallback
from stable_baselines3.common.noise import (
    NormalActionNoise,
    OrnsteinUhlenbeckActionNoise,
)
from stable_baselines3.common.vec_env import DummyVecEnv
import pprint
MODELS = {"a2c": A2C, "ddpg": DDPG, "td3": TD3, "sac": SAC, "ppo": PPO}

MODEL_KWARGS = {x: config.__dict__[f"{x.upper()}_PARAMS"] for x in MODELS.keys()}

NOISE = {
    "normal": NormalActionNoise,
    "ornstein_uhlenbeck": OrnsteinUhlenbeckActionNoise,
}
 
class DRLAgent_SB3:
  def __init__(self,env,run):
    self.env = env
    # self.run = wandb.init(reinit=True,
    #       project = 'finrl-sweeps-sb3',
    #       sync_tensorboard = True,
    #       save_code = True
    #   )
    self.run = run
  def get_model(
      self,
      model_name,
      policy_kwargs=None,
      model_kwargs=None,
      verbose=1,
      seed=None,
  ):
      if model_name not in MODELS:
          raise NotImplementedError("NotImplementedError")

      if model_kwargs is None:
          model_kwargs = MODEL_KWARGS[model_name]

      if "action_noise" in model_kwargs:
          n_actions = self.env.action_space.shape[-1]
          model_kwargs["action_noise"] = NOISE[model_kwargs["action_noise"]](
              mean=np.zeros(n_actions), sigma=0.1 * np.ones(n_actions)
          )
      print(model_kwargs)

      model = MODELS[model_name](
          policy='MlpPolicy',
          env=self.env,
          tensorboard_log=f"runs/{self.run.id}",
          verbose=verbose,
          policy_kwargs=policy_kwargs,
          seed=seed,
          **model_kwargs,
      )
      return model
  
  def train_model(self, model,total_timesteps):
    model = model.learn(
        total_timesteps=total_timesteps,
        callback = WandbCallback(
            gradient_save_freq = 100, model_save_path = f"models/{self.run.id}",
            verbose = 2
        ),
    )
    
    return model
  @staticmethod
  def DRL_prediction_load_from_file(run , model_name, environment,val_or_test='val'):
      if model_name not in MODELS:
          raise NotImplementedError("NotImplementedError, Pass correct model name")
      try:
          # load agent
          model = MODELS[model_name].load(f"models/{run.id}/model.zip") #print_system_info=True
          print("Successfully load model", f"models/{run.id}")
      except BaseException:
          raise ValueError("Fail to load agent!")

      # test on the testing env
      state = environment.reset()
      episode_returns = list()  # the cumulative_return / initial_account
      episode_total_assets = list()
      episode_total_assets.append(environment.initial_total_asset)
      done = False
      while not done:
          action = model.predict(state)[0]
          state, reward, done, _ = environment.step(action)

          total_asset = (
              environment.amount
              + (environment.price_ary[environment.day] * environment.stocks).sum()
          )
          episode_total_assets.append(total_asset)
          episode_return = total_asset / environment.initial_total_asset
          episode_returns.append(episode_return)
    
      def calculate_sharpe(df):
        df['daily_return'] = df['account_value'].pct_change(1)
        if df['daily_return'].std() !=0:
          sharpe = (252**0.5)*df['daily_return'].mean()/ \
              df['daily_return'].std()
          return sharpe
        else:
          return 0

      print("episode_return", episode_return)
      print("Test Finished!")
      sharpe_df = pd.DataFrame(episode_total_assets,columns=['account_value'])
      sharpe = calculate_sharpe(sharpe_df)
      if val_or_test == "val":
        wandb.log({"Val sharpe":sharpe})
      elif val_or_test == "test":
        wandb.log({"Test sharpe":sharpe})
        print(f'Test Sharpe for {run.id} is {sharpe}')
      # run.finish()
      return sharpe, episode_total_assets
from model_wandb import DRLAgent_SB3
def train_agent_env(start_date, end_date, ticker_list, data_source, time_interval, 
          technical_indicator_list, env, model_name, if_vix = True,
          **kwargs):
    
    #fetch data
    DP = DataProcessor(data_source, **kwargs)
    data = DP.download_data(ticker_list, start_date, end_date, time_interval)
    data = DP.clean_data(data)
    data = DP.add_technical_indicator(data, technical_indicator_list)
    if if_vix:
        data = DP.add_vix(data)
    # data.to_csv('train_data.csv')
    # data = pd.read_csv('train_data.csv')
    price_array, tech_array, turbulence_array = DP.df_to_array(data, if_vix)
    env_config = {'price_array':price_array,
              'tech_array':tech_array,
              'turbulence_array':turbulence_array,
              'if_train':True}
    env_instance = env(config=env_config)

    return env_instance

def train(config=None):
    with wandb.init(config=config, sync_tensorboard = True, save_code = True) as run:
      #Get the training environment
      train_env_instance = train_agent_env(TRAIN_START_DATE, TRAIN_END_DATE, ticker_list, data_source, time_interval, 
                            technical_indicator_list, env, model_name)
      config = wandb.config
      #Initialize the training agent
      agent_train = DRLAgent_SB3(train_env_instance,run)
      #For current set of hyperparameters initialize the model
      model = agent_train.get_model(model_name, model_kwargs = config)
      #train the model
      trained_model = agent_train.train_model(model,total_timesteps)
      run_ids[run.id] = run
      print('Training finished!')
      #Log the validation sharpe
      sharpe,val_episode_total_asset = val_or_test(
          VAL_START_DATE, VAL_END_DATE,run,ticker_list, 
          data_source, time_interval, 
          technical_indicator_list, env, model_name
      )
      #Log the testing sharpe
      sharpe,val_episode_total_asset = val_or_test(
          TEST_START_DATE, TEST_END_DATE,run,ticker_list, 
          data_source, time_interval, 
          technical_indicator_list, env, model_name,val_or_test = 'test'
      )
     
def val_or_test(start_date, end_date,run, ticker_list, data_source, time_interval, 
         technical_indicator_list, env, model_name,val_or_test='val', if_vix = True,
         **kwargs):
  
  DP = DataProcessor(data_source, **kwargs)
  data = DP.download_data(ticker_list, start_date, end_date, time_interval)
  data = DP.clean_data(data)
  data = DP.add_technical_indicator(data, technical_indicator_list)
  
  if if_vix:
      data = DP.add_vix(data)
  # if val_or_test == 'val':
  #   data.to_csv('val.csv')
  # elif val_or_test == 'test':
  #   data.to_csv('test.csv')
  # if val_or_test == 'val':
  #   data=pd.read_csv('val.csv')
  # elif val_or_test == 'test':
  #   data = pd.read_csv('test.csv')
  price_array, tech_array, turbulence_array = DP.df_to_array(data, if_vix)
    
  test_env_config = {'price_array':price_array,
          'tech_array':tech_array,
          'turbulence_array':turbulence_array,
          'if_train':False}
  env_instance = env(config=test_env_config)
  
  run_ids[run.id] = run
  sharpe,episode_total_assets = DRLAgent_SB3.DRL_prediction_load_from_file(run,model_name,env_instance,val_or_test)
  
  return sharpe, episode_total_assets
TRAIN_START_DATE = '2009-01-01'
TRAIN_END_DATE = '2019-07-30'

VAL_START_DATE = '2019-08-01'
VAL_END_DATE = '2020-07-30'

TEST_START_DATE = '2020-08-01'
TEST_END_DATE = '2021-10-01'

# ticker_list = config_tickers.DOW_30_TICKER
ticker_list = ['TSLA']
data_source = 'yahoofinance'
time_interval = '1D'
technical_indicator_list = config.TECHNICAL_INDICATORS_LIST
env = StockTradingEnv_numpy
model_name = "a2c"

# PPO_PARAMS = {
#     "n_steps": 2048,
#     "ent_coef": 0.01,
#     "learning_rate": 0.00025,
#     "batch_size": 128,
# }

total_timesteps = 15000
run_ids = {}

count = 30
os.environ['WANDB_AGENT_MAX_INITIAL_FAILURES']= str(count-5)
project_name = 'finrl-sweeps-sb3'
sweep_config = model_params(model_name)

sweep_id = wandb.sweep(sweep_config,project=project_name)
wandb.agent(sweep_id, train, count=count)
run_ids<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11263831258644537, 'fraction_numerical': 0.017721300138312586, 'mean_word_length': 3.2501836884643645, 'pattern_counts': {'":': 40, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '110322', 'license': 'MIT', 'n_tokens_mistral': 4213, 'n_tokens_neox': 4010, 'n_words': 739, 'path': 'tutorials/4-Optimization/FinRL_Weights_and_Biasify_StableBaselines3.ipynb', 'repo_name': 'Neogohan1/FinRL', 'size': '11568'}