FilesExpand file tree

master

/

main.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

69 lines (51 loc) · 2.19 KB

master

/

main.py

Copy path

Top

File metadata and controls

• Code
• Blame

69 lines (51 loc) · 2.19 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

import gym

import json

import datetime as dt

from stable_baselines.common.policies import MlpPolicy

from stable_baselines.common.vec_env import DummyVecEnv

from stable_baselines import PPO2

from env.TradingEnv import TradingEnv

from env.TradingEnv import LOOKFORWARD_WINDOW_SIZE

from env.TradingEnv import LOOKBACK_WINDOW_SIZE

import pandas as pd

TOTAL_TIME_STEPS = 20000

DISPLAY_MODE = 'file' # use mode = 'file' to output files instead of videos

#read the historical stock data

name = input("Please enter the stock historical data file name: ")

asset_name = name[: name.index('_')]

df = pd.read_csv('./data/invest_data/' + name, index_col=0)

df = df.iloc[::-1].reset_index(drop=True)

df = df.sort_values('Date')

train_size = int(len(df) * 0.9)

train_df, test_df = df[0:train_size], df[train_size:len(df)]

test_df = test_df.reset_index(drop=True)

print(train_df.head())

# ARIMA Predictor

from predictors.arima import ArimaPredictor

arima_predictor = ArimaPredictor()

arima_predictor.train(train_df=train_df, column='Close')

fc = arima_predictor.predict(steps=15)

print(fc)

# The algorithms require a vectorized environment to run

train_env = DummyVecEnv([lambda: TradingEnv(train_df)])

test_env = DummyVecEnv([lambda: TradingEnv(test_df)])

model = PPO2(MlpPolicy, train_env, verbose=1)

model.learn(total_timesteps=TOTAL_TIME_STEPS)

model.save(save_path="./saved_model/ppo_{}_{}.pkl".format(asset_name, TOTAL_TIME_STEPS), cloudpickle=True)

obs = train_env.reset()

# back testing on training data

done = False

while not done:

action, _states = model.predict(obs)

obs, rewards, done, info = train_env.step(action)

train_env.render(title=name[:-13], mode=DISPLAY_MODE, filename='LB_{}_LF_{}_{}_{}_train.txt'.

format(LOOKBACK_WINDOW_SIZE, LOOKFORWARD_WINDOW_SIZE, TOTAL_TIME_STEPS, asset_name))

done = False

model.set_env(test_env)

obs = test_env.reset()

# back testing on testing data

while not done:

action, _states = model.predict(obs)

obs, rewards, done, info = test_env.step(action)

test_env.render(title=name[:-13], mode=DISPLAY_MODE, filename='LB_{}_LF_{}_{}_{}_test.txt'.

format(LOOKBACK_WINDOW_SIZE, LOOKFORWARD_WINDOW_SIZE, TOTAL_TIME_STEPS, asset_name))