RonyForAI/Mirage_DB_RL

1import torch2import torch.nn as nn3import torch.optim as optim4import random5import numpy as np6 7 8class Agent(nn.Module):9    def __init__(self, num_tables=3):10        super().__init__()11 12        self.num_tables = num_tables13        self.join_types = 314        self.use_index = 215 16        # total actions = 3 * 3 * 2 = 1817        self.action_size = num_tables * self.join_types * self.use_index18 19        # state size (approx)20        self.state_size = num_tables * 5 + 1  # adjust if needed21 22        self.net = nn.Sequential(23            nn.Linear(self.state_size, 64),24            nn.ReLU(),25            nn.Linear(64, 64),26            nn.ReLU(),27            nn.Linear(64, self.action_size)28        )29 30        self.optimizer = optim.Adam(self.parameters(), lr=0.001)31        self.loss_fn = nn.MSELoss()32 33        # RL params34        self.epsilon = 1.035        self.epsilon_decay = 0.99536        self.epsilon_min = 0.0537        self.gamma = 0.9538 39    def forward(self, x):40        return self.net(x)41 42    # -------- Encode state --------43    def encode_state(self, obs):44        state = []45 46        state.extend(obs.table_rows)        # num_tables values47        state.extend(obs.selectivities)     # num_tables values48        state.extend(obs.has_index)         # num_tables values49 50        # chosen order (pad with -1 to fixed length)51        padded_chosen = obs.chosen_order + [-1] * (self.num_tables - len(obs.chosen_order))52        state.extend(padded_chosen)         # num_tables values53 54        # remaining tables (pad with -1 to fixed length)55        padded_remaining = obs.remaining_tables + [-1] * (self.num_tables - len(obs.remaining_tables))56        state.extend(padded_remaining)      # num_tables values57 58        state.append(obs.current_cost)      # 1 value59 60        # total: num_tables * 5 + 161        return torch.tensor(state, dtype=torch.float32)62 63    # -------- Action encoding --------64    def decode_action(self, action_id):65        table = action_id // 666        rem = action_id % 667        join = rem // 268        index = rem % 269        return table, join, index70 71    # -------- Select action --------72    def select_action(self, obs):73        state = self.encode_state(obs)74 75        if random.random() < self.epsilon:76            action_id = random.randint(0, self.action_size - 1)77        else:78            with torch.no_grad():79                q_values = self.forward(state)80                action_id = torch.argmax(q_values).item()81 82        return self.decode_action(action_id), action_id83 84    # -------- Train step --------85    def train_step(self, state, action, reward, next_state, done):86        q_values = self.forward(state)87        next_q_values = self.forward(next_state)88 89        target = q_values.clone().detach()90        target[action] = reward + (0 if done else self.gamma * torch.max(next_q_values))91 92        loss = self.loss_fn(q_values, target)93 94        self.optimizer.zero_grad()95        loss.backward()96        self.optimizer.step()97 98        # decay epsilon99        if self.epsilon > self.epsilon_min:100            self.epsilon *= self.epsilon_decay