class TradingEnvironment: def __init__(self, stock_df): self.stock_df = stock_df self.current_step = 0 self.total_steps = len(stock_df) - 1 self.reward_range = (0, 1) def reset(self): self.current_step = 0 return self.stock_df.iloc[self.current_step] def step(self, action): self.current_step += 1 done = self.current_step == self.total_steps obs = self.stock_df.iloc[self.current_step] reward = self._get_reward(action) return obs, reward, done def _get_reward(self, action): if action == 0: # 不持有股票 return 0 elif action == 1: # 持有股票 return self.stock_df.iloc[self.current_step]['close'] / self.stock_df.iloc[self.current_step - 1]['close'] - 1 else: raise ValueError("Invalid action, only 0 and 1 are allowed.")
时间: 2024-02-14 09:08:40 浏览: 105
这段代码是一个交易环境类,用于模拟股票交易的过程。其中包括了初始化环境、重置环境、执行动作、获取奖励等方法。具体来说,reset方法用于重置环境,step方法用于执行动作,_get_reward方法用于获取奖励。在执行动作时,可以选择持有股票或不持有股票,持有股票则可以获得当天的收益率,不持有则获得0的奖励。
相关问题
import akshare as ak import numpy as np import pandas as pd import random import matplotlib.pyplot as plt class StockTradingEnv: def __init__(self): self.df = ak.stock_zh_a_daily(symbol='sh000001', adjust="qfq").iloc[::-1] self.observation_space = self.df.shape[1] self.action_space = 3 self.reset() def reset(self): self.current_step = 0 self.total_profit = 0 self.done = False self.state = self.df.iloc[self.current_step].values return self.state def step(self, action): assert self.action_space.contains(action) if action == 0: # 买入 self.buy_stock() elif action == 1: # 卖出 self.sell_stock() else: # 保持不变 pass self.current_step += 1 if self.current_step >= len(self.df) - 1: self.done = True else: self.state = self.df.iloc[self.current_step].values reward = self.get_reward() self.total_profit += reward return self.state, reward, self.done, {} def buy_stock(self): pass def sell_stock(self): pass def get_reward(self): pass class QLearningAgent: def __init__(self, state_size, action_size): self.state_size = state_size self.action_size = action_size self.epsilon = 1.0 self.epsilon_min = 0.01 self.epsilon_decay = 0.995 self.learning_rate = 0.1 self.discount_factor = 0.99 self.q_table = np.zeros((self.state_size, self.action_size)) def act(self, state): if np.random.rand() <= self.epsilon: return random.randrange(self.action_size) else: return np.argmax(self.q_table[state, :]) def learn(self, state, action, reward, next_state, done): target = reward + self.discount_factor * np.max(self.q_table[next_state, :]) self.q_table[state, action] = (1 - self.learning_rate) * self.q_table[state, action] + self.learning_rate * target if self.epsilon > self.epsilon_min: self.epsilon *= self.epsilon_decay env = StockTradingEnv() agent = QLearningAgent(env.observation_space, env.action_space) for episode in range(1000): state = env.reset() done = False while not done: action = agent.act(state) next_state, reward, done, _ = env.step(action) agent.learn(state, action, reward, next_state, done) state = next_state if episode % 10 == 0: print("Episode: %d, Total Profit: %f" % (episode, env.total_profit)) agent.save_model("model-%d.h5" % episode) def plot_profit(env, title): plt.figure(figsize=(12, 6)) plt.plot(env.df.index, env.df.close, label="Price") plt.plot(env.df.index, env.profits, label="Profits") plt.legend() plt.title(title) plt.show() env = StockTradingEnv() agent = QLearningAgent(env.observation_space, env.action_space) agent.load_model("model-100.h5") state = env.reset() done = False while not done: action = agent.act(state) next_state, reward, done, _ = env.step(action) state = next_state plot_profit(env, "QLearning Trading Strategy")优化代码
1. 对于环境类 `StockTradingEnv`,可以考虑将 `buy_stock` 和 `sell_stock` 方法的具体实现写入 `step` 方法中,避免方法数量过多。
2. 可以将 `get_reward` 方法中的具体实现改为直接计算当前持仓的收益。
3. 在循环训练过程中,可以记录每个 episode 的总收益,并将这些数据保存下来,在训练完成后进行可视化分析。
4. 可以添加更多的参数来控制训练过程,比如学习率、衰减系数等。
5. 可以将 QLearningAgent 类中的方法进行整理和封装,提高代码的可读性和可维护性。同时,也可以添加一些对模型进行保存和加载的方法,便于模型的重用和共享。
def train_model(stock_df, agent, num_episodes): for episode in range(num_episodes): obs = stock_df.iloc[0] state = get_state(obs) done = False total_reward = 0 while not done: action = agent.act(state) next_obs = stock_df.iloc[agent.current_step + 1] next_state = get_state(next_obs) reward = get_reward(action, obs, next_obs) total_reward += reward done = agent.current_step == len(stock_df) - 2 agent.learn(state, action, reward, next_state, done) state = next_state obs = next_obs # 输出每个episode的总奖励 print('Episode:', episode, 'Total Reward:', total_reward) # 逐渐降低探索率 agent.set_exploration_rate(agent.exploration_rate * 0.99)修改代码
这段代码中,训练模型时出现了错误,是因为 QLearningAgent 对象没有 current_step 属性,需要在 QLearningAgent 类中实现该属性。你可以在 QLearningAgent 类中添加 current_step 属性,并在 act 方法中更新该属性。例如:
```
class QLearningAgent:
def __init__(self, state_size, action_size):
self.state_size = state_size
self.action_size = action_size
self.q_table = np.zeros((state_size, action_size))
self.exploration_rate = 1.0
self.current_step = 0 # 新增 current_step 属性
def act(self, state):
if np.random.rand() < self.exploration_rate:
return np.random.choice(self.action_size)
q_values = self.q_table[state]
return np.argmax(q_values)
def learn(self, state, action, reward, next_state, done):
# 更新 Q 表
q_next = self.q_table[next_state]
if done:
q_next = np.zeros(self.action_size)
td_target = reward + GAMMA * np.max(q_next)
td_error = td_target - self.q_table[state, action]
self.q_table[state, action] += ALPHA * td_error
# 更新探索率和当前步数
self.current_step += 1
self.exploration_rate *= EXPLORATION_DECAY
self.exploration_rate = max(EXPLORATION_MIN, self.exploration_rate)
```
在修改后的代码中,我们在 QLearningAgent 类中新增了 current_step 属性,并在 act 方法和 learn 方法中更新该属性。最后,在训练模型时,我们可以使用 QLearningAgent 对象的 current_step 属性来获取当前步数,而不会再出现属性错误。
阅读全文