checkpoints

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 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), sell at the next candle's open and compute reward.
+          - If in a position and action is SELL (0), close the position at the next candle's open.
-          - Otherwise, no trade is executed.
+        Returns: (current_state, reward, next_state, done)
-        Returns: (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.
+    For each epoch, run through the entire historical episode.
-    At each step, choose an action using ε‑greedy policy, simulate a trade,
+    At each step, pick an action (using ε-greedy), simulate a trade, store the experience,
-    store the experience (state, action, reward, next_state, done), and update the model.
+    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
 
@@ -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