amayasanduni/plan-generator

1import torch
2import torch.nn as nn
3import torch.optim as optim
4import numpy as np
5import random
6from collections import deque
7
8# --- 1. Model Architecture (Must match main.py) ---
9class DQN(nn.Module):
10    def __init__(self, state_size, action_size):
11        super(DQN, self).__init__()
12        self.fc1 = nn.Linear(state_size, 64)
13        self.fc2 = nn.Linear(64, 64)
14        self.fc3 = nn.Linear(64, action_size)
15
16    def forward(self, x):
17        x = torch.relu(self.fc1(x))
18        x = torch.relu(self.fc2(x))
19        return self.fc3(x)
20
21# --- 2. Replay Memory ---
22class ReplayMemory:
23    def __init__(self, capacity):
24        self.memory = deque(maxlen=capacity)
25    def push(self, transition):
26        self.memory.append(transition)
27    def sample(self, batch_size):
28        return random.sample(self.memory, batch_size)
29    def __len__(self):
30        return len(self.memory)
31
32# --- 3. The Balanced Environment ---
33class StudentEnv:
34    def __init__(self, scores, user_level_idx):
35        self.scores = np.array(scores, dtype=float)
36        self.user_level_idx = user_level_idx
37        self.thresholds = np.array([50, 50, 50])
38
39    def get_state(self):
40        weakness = (self.scores < 45).astype(int)
41        # 3 scores + 3 weakness + 1 level = 7
42        return np.concatenate((self.scores, weakness, [float(self.user_level_idx)]))
43
44    def step(self, action):
45        # Identify the priority order of scores
46        # The AI should get the most reward for the lowest score, 
47        # but some reward for the second lowest.
48        sorted_score_indices = np.argsort(self.scores)
49        primary_weakness = sorted_score_indices[0]
50        secondary_weakness = sorted_score_indices[1]
51
52        reward = 0
53        if action == primary_weakness:
54            reward = 25  # Highest reward for major gap
55        elif action == secondary_weakness:
56            reward = 10  # Moderate reward for second gap
57        elif action == 3: # Mixed
58            reward = 5
59        else:
60            reward = -10 # Penalty for ignoring weaknesses
61
62        # Simulate Improvement
63        if action < 3:
64            self.scores[action] += random.randint(5, 10)
65        else:
66            self.scores += random.randint(2, 4)
67
68        self.scores = np.clip(self.scores, 0, 100)
69        done = np.all(self.scores >= 55)
70        return self.get_state(), reward, done
71
72# --- 4. The Agent ---
73class DQNAgent:
74    def __init__(self, state_size, action_size):
75        self.state_size = state_size
76        self.action_size = action_size
77        self.memory = ReplayMemory(5000)
78        self.gamma = 0.95
79        self.epsilon = 1.0
80        self.epsilon_min = 0.01
81        self.epsilon_decay = 0.996 
82        self.model = DQN(state_size, action_size)
83        self.optimizer = optim.Adam(self.model.parameters(), lr=0.001)
84        self.criterion = nn.MSELoss()
85
86    def act(self, state):
87        if np.random.rand() < self.epsilon:
88            return random.randrange(self.action_size)
89        state_t = torch.FloatTensor(state).unsqueeze(0)
90        with torch.no_grad():
91            return torch.argmax(self.model(state_t)).item()
92
93    def replay(self, batch_size):
94        if len(self.memory) < batch_size: return
95        minibatch = self.memory.sample(batch_size)
96        for state, action, reward, next_state in minibatch:
97            target = reward + self.gamma * torch.max(self.model(torch.FloatTensor(next_state).unsqueeze(0)))
98            current = self.model(torch.FloatTensor(state).unsqueeze(0))[0][action]
99            loss = self.criterion(current, target)
100            self.optimizer.zero_grad(); loss.backward(); self.optimizer.step()
101        if self.epsilon > self.epsilon_min: self.epsilon *= self.epsilon_decay
102
103# --- 5. Training Execution ---
104def run_training():
105    state_size = 7
106    action_size = 4
107    agent = DQNAgent(state_size, action_size)
108    episodes = 800 # Higher episodes for better ranking accuracy
109
110    
111
112    print("Training Balanced Ranking Model...")
113    for e in range(episodes):
114        level = random.randint(0, 2)
115        start_scores = [random.randint(10, 50) for _ in range(3)]
116        env = StudentEnv(start_scores, level)
117        state = env.get_state()
118        
119        for _ in range(15):
120            action = agent.act(state)
121            next_state, reward, done = env.step(action)
122            agent.memory.push((state, action, reward, next_state))
123            state = next_state
124            agent.replay(32)
125            if done: break
126            
127        if (e + 1) % 100 == 0:
128            print(f"Episode: {e+1}/{episodes}, Exploration: {agent.epsilon:.2f}")
129
130    torch.save(agent.model.state_dict(), "dqn_roadmap_model.pth")
131    print("Success! New model trained with 7-element state and ranking logic.")
132
133if __name__ == "__main__":
134    run_training()