demeter

  1# agent/Marl/Bandit.py
  2
  3import numpy as np
  4import pickle
  5import os
  6
  7class ContextualBandit:
  8    def __init__(self, n_actions=15, feature_dim=519):
  9        """
 10        LinGreedy Implementation (Pure Exploitation).
 11        We removed 'alpha' because we do not want to explore.
 12        """
 13        self.n_actions = n_actions
 14        self.d = feature_dim
 15        
 16        # A: Covariance Matrix (Used for Ridge Regression learning)
 17        self.A = [np.identity(self.d) for _ in range(self.n_actions)]
 18        
 19        # b: Reward Vector
 20        self.b = [np.zeros(self.d) for _ in range(self.n_actions)]
 21        
 22        # theta: Weight vectors for each action (optional, computed on-the-fly)
 23        self.theta = [np.zeros(self.d) for _ in range(self.n_actions)]
 24        
 25        self.file_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "model_bandit_greedy.pkl")
 26        self.load()  # Now calls with no arguments
 27
 28    def select_action(self, context_vector):
 29        """
 30        Returns: (action_index, debug_info)
 31        Strictly picks the action with the highest PREDICTED reward.
 32        """
 33        # Ensure context is flat (1D array)
 34        context_vector = np.array(context_vector).reshape(-1)
 35        
 36        predicted_rewards = np.zeros(self.n_actions)
 37        confidences = np.zeros(self.n_actions)
 38        
 39        for a in range(self.n_actions):
 40            # 1. Calculate Mean Estimate (theta)
 41            # theta = A^-1 * b
 42            try:
 43                theta = np.linalg.solve(self.A[a], self.b[a])
 44            except np.linalg.LinAlgError:
 45                # Fallback for singular matrix (rare with Identity init)
 46                theta = np.zeros(self.d)
 47            
 48            # 2. Expected Reward (Dot Product)
 49            # This is the "Best Guess" for how good this action is.
 50
 51        #    print(f"Action {a}: Theta shape: {theta.shape}, Context shape: {context_vector.shape}")
 52            pred = theta.dot(context_vector)
 53            predicted_rewards[a] = pred
 54            
 55            # 3. Calculate Confidence (Optional, for UI only)
 56            # We calculate variance just to show the user "How sure are we?"
 57            # But we do NOT add this to the score.
 58            try:
 59                variance = context_vector.dot(np.linalg.solve(self.A[a], context_vector))
 60                confidences[a] = 1.0 / (1.0 + variance) # Simple confidence score (0-1)
 61            except np.linalg.LinAlgError:
 62                confidences[a] = 0.0
 63
 64        # 🟢 PURE EXPLOITATION: Pick max predicted reward
 65        chosen_action = np.argmax(predicted_rewards)
 66        
 67        return chosen_action, {
 68            "scores": predicted_rewards.tolist(), 
 69            "confidences": confidences.tolist()
 70        }
 71
 72    def update(self, context_vector, action_idx, reward):
 73        """
 74        Online Learning: The AI still gets smarter with every feedback.
 75        
 76        Args:
 77            context_vector: The feature vector (515-dim)
 78            action_idx: Which action was taken
 79            reward: The reward received
 80        """
 81        # Ensure types are correct
 82        action_idx = int(action_idx)
 83        reward = float(reward)
 84        
 85        # Ensure context is flat (1D array)
 86        context_vector = np.array(context_vector).reshape(-1)
 87        
 88        # Update the regression model for the chosen arm
 89        self.A[action_idx] += np.outer(context_vector, context_vector)
 90        self.b[action_idx] += reward * context_vector
 91        
 92        # Update theta (optional, can be computed on-the-fly in select_action)
 93        try:
 94            self.theta[action_idx] = np.linalg.solve(self.A[action_idx], self.b[action_idx])
 95        except np.linalg.LinAlgError:
 96            # Use pseudoinverse if solve fails
 97            self.theta[action_idx] = np.linalg.pinv(self.A[action_idx]).dot(self.b[action_idx])
 98
 99    def save(self):
100        """Save the model weights to disk"""
101        try:
102            with open(self.file_path, 'wb') as f:
103                pickle.dump({
104                    'A': self.A, 
105                    'b': self.b,
106                    'theta': self.theta
107                }, f)
108            print(f"💾 Model saved to {self.file_path}")
109        except Exception as e:
110            print(f"❌ Error saving model: {e}")
111
112    def load(self):  # <--- FIXED: No filepath argument, uses self.file_path
113        """Loads weights from disk if they exist."""
114        if not os.path.exists(self.file_path):
115            print(f"ℹ️ No saved model found at {self.file_path}. Starting fresh.")
116            return False
117            
118        try:
119            with open(self.file_path, 'rb') as f:
120                state = pickle.load(f)
121            
122            # Restore state
123            self.A = state['A']
124            self.b = state['b']
125            self.theta = state['theta']
126            print(f"✅ Loaded bandit model from {self.file_path}")
127            return True
128        except Exception as e:
129            print(f"⚠️ Error loading model: {e}. Starting fresh.")
130            return False