diff --git a/crypto/brian/index.py b/crypto/brian/index.py index 4000377..98d8990 100644 --- a/crypto/brian/index.py +++ b/crypto/brian/index.py @@ -14,12 +14,19 @@ import torch.nn as nn import torch.optim as optim import numpy as np from collections import deque +from datetime import datetime + +# --- Directories for saving models --- +LAST_DIR = os.path.join("models", "last") +BEST_DIR = os.path.join("models", "best") +os.makedirs(LAST_DIR, exist_ok=True) +os.makedirs(BEST_DIR, exist_ok=True) + +CACHE_FILE = "candles_cache.json" # ------------------------------------- # Utility functions for caching candles to file # ------------------------------------- -CACHE_FILE = "candles_cache.json" - def load_candles_cache(filename): if os.path.exists(filename): try: @@ -38,6 +45,81 @@ def save_candles_cache(filename, candles): except Exception as e: print("Error saving cache file:", e) +# ------------------------------------- +# Functions for handling checkpoints +# ------------------------------------- +def maintain_checkpoint_directory(directory, max_files=10): + """Keep only the most recent max_files in a given directory based on modification time.""" + files = os.listdir(directory) + if len(files) > max_files: + full_paths = [os.path.join(directory, f) for f in files] + full_paths.sort(key=lambda x: os.path.getmtime(x)) + # Remove the oldest files + for f in full_paths[: len(files) - max_files]: + os.remove(f) + +def get_best_models(directory): + """Return a list of (reward, filename) for files in the best folder. + Expecting filenames like: best_{reward:.4f}_epoch_{epoch}_{timestamp}.pt""" + best_files = [] + for file in os.listdir(directory): + parts = file.split("_") + try: + # parts[1] should be reward + r = float(parts[1]) + best_files.append((r, file)) + except Exception: + continue + return best_files + +def save_checkpoint(model, epoch, reward, last_dir=LAST_DIR, best_dir=BEST_DIR): + """Save the model state always to the last_dir and conditionally to best_dir if reward is high enough.""" + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + last_filename = f"model_last_epoch_{epoch}_{timestamp}.pt" + last_path = os.path.join(last_dir, last_filename) + torch.save({ + "epoch": epoch, + "reward": reward, + "model_state_dict": model.state_dict() + }, last_path) + # Keep only last 10 models in last_dir. + maintain_checkpoint_directory(last_dir, max_files=10) + + # Check the best folder – if fewer than 10, simply add; + # Otherwise, add only if reward is higher than the lowest reward in best. + best_models = get_best_models(best_dir) + add_to_best = False + if len(best_models) < 10: + add_to_best = True + else: + min_reward, min_file = min(best_models, key=lambda x: x[0]) + if reward > min_reward: + add_to_best = True + # Remove the worst checkpoint. + os.remove(os.path.join(best_dir, min_file)) + if add_to_best: + best_filename = f"best_{reward:.4f}_epoch_{epoch}_{timestamp}.pt" + best_path = os.path.join(best_dir, best_filename) + torch.save({ + "epoch": epoch, + "reward": reward, + "model_state_dict": model.state_dict() + }, best_path) + maintain_checkpoint_directory(best_dir, max_files=10) + print(f"Saved checkpoint for epoch {epoch} with reward {reward:.4f}") + +def load_best_checkpoint(model, best_dir=BEST_DIR): + """Load the best checkpoint (with highest reward) from the best directory if available.""" + best_models = get_best_models(best_dir) + if not best_models: + return None + best_reward, best_file = max(best_models, key=lambda x: x[0]) + path = os.path.join(best_dir, best_file) + print(f"Loading best model from checkpoint: {best_file} with reward {best_reward:.4f}") + checkpoint = torch.load(path) + model.load_state_dict(checkpoint["model_state_dict"]) + return checkpoint + # ------------------------------------- # Neural Network Architecture Definition # ------------------------------------- @@ -79,7 +161,7 @@ def compute_indicators(candle, additional_data): """ Combine OHLCV candle data with extra indicator information. Base features: open, high, low, close, volume. - Additional channels (e.g. simulated sentiment) are appended. + Additional channels (e.g., simulated sentiment) are appended. """ features = [ candle.get('open', 0.0), @@ -105,7 +187,7 @@ class ContinuousRLAgent: self.gamma = gamma def act(self, state, epsilon=0.1): - # ε-greedy: with probability epsilon take a random action + # ε-greedy: choose random action with probability epsilon. if np.random.rand() < epsilon: return np.random.randint(0, 3) state_tensor = torch.from_numpy(np.array(state, dtype=np.float32)).unsqueeze(0) @@ -115,12 +197,12 @@ class ContinuousRLAgent: return action def train_step(self): - # Only train if we have enough samples + # Only train if we have enough samples. if len(self.replay_buffer) < self.batch_size: return - # Convert lists to numpy arrays in one shot for performance batch = self.replay_buffer.sample(self.batch_size) + # Unpack the batch; each experience is (state, action, reward, next_state, done) states, actions, rewards, next_states, dones = zip(*batch) states_tensor = torch.from_numpy(np.array(states, dtype=np.float32)) actions_tensor = torch.tensor(actions, dtype=torch.int64) @@ -128,15 +210,12 @@ class ContinuousRLAgent: next_states_tensor = torch.from_numpy(np.array(next_states, dtype=np.float32)) dones_tensor = torch.tensor(dones, dtype=torch.float32).unsqueeze(1) - # Current Q-value for the chosen actions Q_values = self.model(states_tensor) current_Q = Q_values.gather(1, actions_tensor.unsqueeze(1)) - with torch.no_grad(): next_Q_values = self.model(next_states_tensor) max_next_Q = next_Q_values.max(1)[0].unsqueeze(1) target = rewards_tensor + self.gamma * max_next_Q * (1.0 - dones_tensor) - loss = self.loss_fn(current_Q, target) self.optimizer.zero_grad() loss.backward() @@ -148,7 +227,7 @@ class ContinuousRLAgent: async def fetch_historical_data(exchange, symbol, timeframe, since, end_time, batch_size=500): """ Fetch historical OHLCV data for the given symbol and timeframe. - The 'since' and 'end_time' parameters are in milliseconds. + "since" and "end_time" are in milliseconds. """ candles = [] since_ms = since @@ -181,7 +260,6 @@ async def get_cached_or_fetch_data(exchange, symbol, timeframe, since, end_time, cached_candles = load_candles_cache(cache_file) if cached_candles: last_ts = cached_candles[-1]['timestamp'] - # If the cached candles do not extend to 'end_time', fetch new ones. if last_ts < end_time: print("Fetching new candles to update cache...") new_candles = await fetch_historical_data(exchange, symbol, timeframe, last_ts + 1, end_time, batch_size) @@ -209,7 +287,6 @@ class BacktestEnvironment: def get_state(self, index): candle = self.candles[index] - # Simulate additional sentiment features. sentiment = { 'sentiment_score': np.random.rand(), 'news_volume': np.random.rand(), @@ -220,10 +297,9 @@ class BacktestEnvironment: def step(self, action): """ Simulate a trading step. - - If not in a position and action is BUY (2), buy at the next candle's open. - - If in a position and action is SELL (0), sell at the next candle's open and compute reward. - - Otherwise, no trade is executed. - Returns: (state, reward, next_state, done) + - If not in a position and action is BUY (2), enter a long position at the next candle's open. + - If in a position and action is SELL (0), close the position at the next candle's open. + Returns: (current_state, reward, next_state, done) """ if self.current_index >= len(self.candles) - 1: return self.get_state(self.current_index), 0.0, None, True @@ -237,14 +313,15 @@ class BacktestEnvironment: # Action mapping: 0 -> SELL, 1 -> HOLD, 2 -> BUY. if self.position is None: - if action == 2: # BUY: enter long position at next candle's open. + if action == 2: # BUY signal: entry_price = next_candle['open'] self.position = {'entry_price': entry_price, 'entry_index': self.current_index} else: - if action == 0: # SELL: close long position. + if action == 0: # SELL signal: sell_price = next_candle['open'] reward = sell_price - self.position['entry_price'] self.position = None + self.current_index = next_index done = (self.current_index >= len(self.candles) - 1) return current_state, reward, next_state, done @@ -254,11 +331,11 @@ class BacktestEnvironment: # ------------------------------------- def train_on_historical_data(env, rl_agent, num_epochs=10, epsilon=0.1): """ - For each epoch, run through the entire historical data. - At each step, choose an action using ε‑greedy policy, simulate a trade, - store the experience (state, action, reward, next_state, done), and update the model. + For each epoch, run through the entire historical episode. + At each step, pick an action (using ε-greedy), simulate a trade, store the experience, + and update the model. Then log the cumulative reward and save checkpoints. """ - for epoch in range(num_epochs): + for epoch in range(1, num_epochs + 1): state = env.reset() done = False total_reward = 0.0 @@ -269,12 +346,14 @@ def train_on_historical_data(env, rl_agent, num_epochs=10, epsilon=0.1): state, reward, next_state, done = env.step(action) if next_state is None: next_state = np.zeros_like(prev_state) - # Store the experience including the action taken. + # Save the experience (state, action, reward, next_state, done) rl_agent.replay_buffer.add((prev_state, action, reward, next_state, done)) rl_agent.train_step() total_reward += reward steps += 1 - print(f"Epoch {epoch+1}/{num_epochs} completed, total reward: {total_reward:.4f} over {steps} steps.") + print(f"Epoch {epoch}/{num_epochs} completed, total reward: {total_reward:.4f} over {steps} steps.") + # Save a checkpoint after the epoch. + save_checkpoint(rl_agent.model, epoch, total_reward, LAST_DIR, BEST_DIR) # ------------------------------------- # Main Asynchronous Function for Backtest Training @@ -285,7 +364,7 @@ async def main_backtest(): timeframe = '1m' now = int(time.time() * 1000) one_day_ms = 24 * 60 * 60 * 1000 - # Fetch a 1-day period from 2 days ago until 1 day ago. + # For example, fetch a 1-day period from 2 days ago until 1 day ago. since = now - one_day_ms * 2 end_time = now - one_day_ms @@ -297,7 +376,7 @@ async def main_backtest(): 'secret': mexc_api_secret, 'enableRateLimit': True, }) - + print("Fetching historical data...") candles = await get_cached_or_fetch_data(exchange, symbol, timeframe, since, end_time) if not candles: @@ -305,27 +384,30 @@ async def main_backtest(): await exchange.close() return - # Save/Update cache file. + # Save updated cache. save_candles_cache(CACHE_FILE, candles) - # Initialize the backtest environment with the candles. + # Initialize backtest environment. env = BacktestEnvironment(candles) - - # Model dimensions: 5 base OHLCV features + 3 simulated sentiment features = 8. + + # Model dimensions: 5 (OHLCV) + 3 (sentiment) = 8. input_dim = 8 hidden_dim = 128 - output_dim = 3 # SELL, HOLD, BUY + output_dim = 3 # SELL, HOLD, BUY. model = TradingModel(input_dim, hidden_dim, output_dim) optimizer = optim.Adam(model.parameters(), lr=1e-4) replay_buffer = ReplayBuffer(capacity=10000) rl_agent = ContinuousRLAgent(model, optimizer, replay_buffer, batch_size=32, gamma=0.99) - + + # At training start, try loading a best checkpoint (if available). + load_best_checkpoint(model, BEST_DIR) + # Run training over historical data. - num_epochs = 10 # Adjust as needed. + num_epochs = 10 # Change as needed. train_on_historical_data(env, rl_agent, num_epochs=num_epochs, epsilon=0.1) - - # Optionally, perform a final test run (without exploration) to check cumulative profit. + + # Final simulation (without exploration) to check cumulative profit. state = env.reset() done = False cumulative_reward = 0.0