Spaces:
Sleeping
Sleeping
Viswanath129
Complete overhaul to fix phase 1: consolidated server package, updated Dockerfile, and fixed imports
3fade06 | import torch | |
| import torch.nn as nn | |
| import torch.optim as optim | |
| import numpy as np | |
| import random | |
| from collections import deque | |
| import os | |
| class DQN(nn.Module): | |
| def __init__(self, state_size, action_size): | |
| super(DQN, self).__init__() | |
| self.fc1 = nn.Linear(state_size, 64) | |
| self.fc2 = nn.Linear(64, 64) | |
| self.fc3 = nn.Linear(64, action_size) | |
| def forward(self, x): | |
| x = torch.relu(self.fc1(x)) | |
| x = torch.relu(self.fc2(x)) | |
| return self.fc3(x) | |
| class DQNAgent: | |
| def __init__(self, state_size=5, action_size=4): | |
| self.state_size = state_size | |
| self.action_size = action_size | |
| self.memory = deque(maxlen=2000) | |
| self.gamma = 0.95 | |
| self.epsilon = 1.0 | |
| self.epsilon_min = 0.01 | |
| self.epsilon_decay = 0.995 | |
| self.learning_rate = 0.001 | |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| self.model = DQN(state_size, action_size).to(self.device) | |
| self.optimizer = optim.Adam(self.model.parameters(), lr=self.learning_rate) | |
| self.criterion = nn.MSELoss() | |
| def get_action(self, state, explore=True): | |
| if explore and np.random.rand() <= self.epsilon: | |
| return random.randrange(self.action_size) | |
| state_t = torch.FloatTensor(state).to(self.device).unsqueeze(0) | |
| with torch.no_grad(): | |
| act_values = self.model(state_t) | |
| return torch.argmax(act_values[0]).item() | |
| def remember(self, state, action, reward, next_state, done): | |
| self.memory.append((state, action, reward, next_state, done)) | |
| def replay(self, batch_size): | |
| if len(self.memory) < batch_size: | |
| return | |
| minibatch = random.sample(self.memory, batch_size) | |
| for state, action, reward, next_state, done in minibatch: | |
| state_t = torch.FloatTensor(state).to(self.device) | |
| next_state_t = torch.FloatTensor(next_state).to(self.device) | |
| target = reward | |
| if not done: | |
| target = reward + self.gamma * torch.max(self.model(next_state_t.unsqueeze(0))[0]).item() | |
| target_f = self.model(state_t.unsqueeze(0)) | |
| target_f[0][action] = target | |
| self.optimizer.zero_grad() | |
| output = self.model(state_t.unsqueeze(0)) | |
| loss = self.criterion(output, target_f) | |
| loss.backward() | |
| self.optimizer.step() | |
| if self.epsilon > self.epsilon_min: | |
| self.epsilon *= self.epsilon_decay | |
| def save(self, path): | |
| os.makedirs(os.path.dirname(path), exist_ok=True) | |
| torch.save(self.model.state_dict(), path) | |
| def load(self, path): | |
| if os.path.exists(path): | |
| self.model.load_state_dict(torch.load(path, map_location=self.device)) | |
| self.epsilon = self.epsilon_min | |