Upload 2 files

dqn.py

@@ -0,0 +1,93 @@
+import numpy as np
+import pandas as pd
+import random
+from sklearn.utils import shuffle
+import torch
+import torch.nn as nn
+import torch.autograd as autograd
+from torchcontrib.optim import SWA
+from collections import deque
+
+from preprocess import *
+
+class DQN(nn.Module):
+    
+    def __init__(self, input_dim, output_dim):
+        super(DQN, self).__init__()
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        self.fc = nn.Sequential( \
+            nn.Linear(self.input_dim[0], 32), \
+            nn.ReLU(), \
+            nn.Linear(32, self.output_dim))
+
+    def forward(self, state):
+        return self.fc(state)
+    
+class DQNAgent:
+
+    def __init__(self, input_dim, dataset,
+                 learning_rate=3e-4, 
+                 gamma=0.99,
+                 buffer=None,
+                 buffer_size=10000, 
+                 tau=0.999,
+                 swa=False,
+                 pre_trained_model=None):
+        self.learning_rate = learning_rate
+        self.gamma = gamma
+        self.tau = tau
+        self.model = DQN(input_dim, 1)
+        if pre_trained_model:
+            self.model = pre_trained_model
+        base_opt = torch.optim.Adam(self.model.parameters())
+        self.swa = swa
+        self.dataset=dataset
+        self.MSE_loss = nn.MSELoss()
+        self.replay_buffer = buffer
+        if swa:
+          self.optimizer = SWA(base_opt, swa_start=10, swa_freq=5, swa_lr=0.05)
+        else:
+          self.optimizer = base_opt
+
+    def get_action(self, state, dataset=None):
+        if dataset is None:
+            dataset = self.dataset
+        inputs = get_multiple_model_inputs(state, state.remaining, dataset)
+        model_inputs = autograd.Variable(torch.from_numpy(inputs).float().unsqueeze(0))
+        expected_returns = self.model.forward(model_inputs)
+        value, index = expected_returns.max(1)
+        return state.remaining[index[0]]
+
+    def compute_loss(self, batch, dataset, verbose=False):
+        states, actions, rewards, next_states, dones = batch
+        model_inputs = np.array([get_model_inputs(states[i], actions[i], dataset)\
+            for i in range(len(states))])
+        model_inputs = torch.FloatTensor(model_inputs)
+
+        rewards = torch.FloatTensor(rewards)
+        dones = torch.FloatTensor(dones)
+
+        curr_Q = self.model.forward(model_inputs)
+        model_inputs = np.array([get_model_inputs(next_states[i], actions[i], dataset) \
+            for i in range(len(next_states))])
+        model_inputs = torch.FloatTensor(model_inputs)
+        next_Q = self.model.forward(model_inputs)
+        max_next_Q = torch.max(next_Q, 1)[0]
+        expected_Q = rewards.squeeze(1) + (1 - dones) * self.gamma * max_next_Q
+
+        if verbose:
+            print(curr_Q, expected_Q)
+        loss = self.MSE_loss(curr_Q.squeeze(0), expected_Q.detach())
+        return loss
+
+    def update(self, batch_size, verbose=False):
+        batch = self.replay_buffer.sample(batch_size)
+        loss = self.compute_loss(batch, self.dataset, verbose)
+        train_loss = loss.float()
+        self.optimizer.zero_grad()
+        loss.backward()
+        self.optimizer.step()
+        if self.swa:
+            self.optimizer.swap_swa_sgd()
+        return train_loss

mdp.py

@@ -0,0 +1,86 @@
+# State and Buffer Classes
+
+import numpy as np
+import pandas as pd
+import random
+from sklearn.utils import shuffle
+import torch
+import torch.nn as nn
+import torch.autograd as autograd
+from torchcontrib.optim import SWA
+from collections import deque
+
+from preprocess import *
+
+def compute_reward(t, relevance):
+    """
+    Reward function for MDP
+    """
+    if t == 0:
+        return 0
+    return relevance / np.log2(t + 1)
+
+class State:
+
+    def __init__(self, t, query, remaining):
+        self.t = t
+        self.qid = query #useful for sorting buffer
+        self.remaining = remaining
+
+    def pop(self):
+        return self.remaining.pop()
+
+    def initial(self):
+        return self.t == 0
+
+    def terminal(self):
+        return len(self.remaining) == 0
+
+class BasicBuffer:
+
+    def __init__(self, max_size):
+        self.max_size = max_size
+        self.buffer = deque(maxlen=max_size)
+
+    def push(self, state, action, reward, next_state, done):
+        experience = (state, action, np.array([reward]), next_state, done)
+        self.buffer.append(experience)
+
+    def push_batch(self, df, n):
+        for i in range(n):
+            random_qid = random.choice(list(df["qid"]))
+            filtered_df = df.loc[df["qid"] == int(random_qid)].reset_index()
+            row_order = [x for x in range(len(filtered_df))]
+            X = [x[1]["doc_id"] for x in filtered_df.iterrows()]
+            random.shuffle(row_order)
+            for t,r in enumerate(row_order):
+                cur_row = filtered_df.iloc[r]
+                old_state = State(t, cur_row["qid"], X[:])
+                action = cur_row["doc_id"]
+                new_state = State(t+1, cur_row["qid"], X[:])
+                reward = compute_reward(t+1, cur_row["rank"])
+                self.push(old_state, action, reward, new_state, t+1 == len(row_order))
+                filtered_df.drop(filtered_df.index[[r]])
+
+    def sample(self, batch_size):
+        state_batch = []
+        action_batch = []
+        reward_batch = []
+        next_state_batch = []
+        done_batch = []
+
+        batch = random.sample(self.buffer, batch_size)
+
+        for experience in batch:
+            state, action, reward, next_state, done = experience
+            state_batch.append(state)
+            action_batch.append(action)
+            reward_batch.append(reward)
+            next_state_batch.append(next_state)
+            done_batch.append(done)
+
+        return (state_batch, action_batch, reward_batch, 
+            next_state_batch, done_batch)
+
+    def __len__(self):
+        return len(self.buffer)