A 3-stock trading environment. State: vector of size 7 (n_stock * 2 + 1) - # shares of stock 1 owned - # shares of stock 2 owned - # shares of stock 3 owned - price of stock 1 (using daily close price) - price of stock 2 - price of stock 3 - cash owned (can be used t
| 139 | |
| 140 | |
| 141 | class MultiStockEnv: |
| 142 | """ |
| 143 | A 3-stock trading environment. |
| 144 | State: vector of size 7 (n_stock * 2 + 1) |
| 145 | - # shares of stock 1 owned |
| 146 | - # shares of stock 2 owned |
| 147 | - # shares of stock 3 owned |
| 148 | - price of stock 1 (using daily close price) |
| 149 | - price of stock 2 |
| 150 | - price of stock 3 |
| 151 | - cash owned (can be used to purchase more stocks) |
| 152 | Action: categorical variable with 27 (3^3) possibilities |
| 153 | - for each stock, you can: |
| 154 | - 0 = sell |
| 155 | - 1 = hold |
| 156 | - 2 = buy |
| 157 | """ |
| 158 | def __init__(self, data, initial_investment=20000): |
| 159 | # data |
| 160 | self.stock_price_history = data |
| 161 | self.n_step, self.n_stock = self.stock_price_history.shape |
| 162 | |
| 163 | # instance attributes |
| 164 | self.initial_investment = initial_investment |
| 165 | self.cur_step = None |
| 166 | self.stock_owned = None |
| 167 | self.stock_price = None |
| 168 | self.cash_in_hand = None |
| 169 | |
| 170 | self.action_space = np.arange(3**self.n_stock) |
| 171 | |
| 172 | # action permutations |
| 173 | # returns a nested list with elements like: |
| 174 | # [0,0,0] |
| 175 | # [0,0,1] |
| 176 | # [0,0,2] |
| 177 | # [0,1,0] |
| 178 | # [0,1,1] |
| 179 | # etc. |
| 180 | # 0 = sell |
| 181 | # 1 = hold |
| 182 | # 2 = buy |
| 183 | self.action_list = list(map(list, itertools.product([0, 1, 2], repeat=self.n_stock))) |
| 184 | |
| 185 | # calculate size of state |
| 186 | self.state_dim = self.n_stock * 2 + 1 |
| 187 | |
| 188 | self.reset() |
| 189 | |
| 190 | |
| 191 | def reset(self): |
| 192 | self.cur_step = 0 |
| 193 | self.stock_owned = np.zeros(self.n_stock) |
| 194 | self.stock_price = self.stock_price_history[self.cur_step] |
| 195 | self.cash_in_hand = self.initial_investment |
| 196 | return self._get_obs() |
| 197 | |
| 198 |