| 282 | |
| 283 | |
| 284 | class DQNAgent(object): |
| 285 | def __init__(self, state_size, action_size): |
| 286 | self.state_size = state_size |
| 287 | self.action_size = action_size |
| 288 | self.memory = ReplayBuffer(state_size, action_size, size=500) |
| 289 | self.gamma = 0.95 # discount rate |
| 290 | self.epsilon = 1.0 # exploration rate |
| 291 | self.epsilon_min = 0.01 |
| 292 | self.epsilon_decay = 0.995 |
| 293 | self.model = MLP(state_size, action_size) |
| 294 | |
| 295 | # Loss and optimizer |
| 296 | self.criterion = nn.MSELoss() |
| 297 | self.optimizer = torch.optim.Adam(self.model.parameters()) |
| 298 | |
| 299 | |
| 300 | def update_replay_memory(self, state, action, reward, next_state, done): |
| 301 | self.memory.store(state, action, reward, next_state, done) |
| 302 | |
| 303 | |
| 304 | def act(self, state): |
| 305 | if np.random.rand() <= self.epsilon: |
| 306 | return np.random.choice(self.action_size) |
| 307 | act_values = predict(self.model, state) |
| 308 | return np.argmax(act_values[0]) # returns action |
| 309 | |
| 310 | |
| 311 | def replay(self, batch_size=32): |
| 312 | # first check if replay buffer contains enough data |
| 313 | if self.memory.size < batch_size: |
| 314 | return |
| 315 | |
| 316 | # sample a batch of data from the replay memory |
| 317 | minibatch = self.memory.sample_batch(batch_size) |
| 318 | states = minibatch['s'] |
| 319 | actions = minibatch['a'] |
| 320 | rewards = minibatch['r'] |
| 321 | next_states = minibatch['s2'] |
| 322 | done = minibatch['d'] |
| 323 | |
| 324 | # Calculate the target: Q(s',a) |
| 325 | target = rewards + (1 - done) * self.gamma * np.amax(predict(self.model, next_states), axis=1) |
| 326 | |
| 327 | # With the PyTorch API, it is simplest to have the target be the |
| 328 | # same shape as the predictions. |
| 329 | # However, we only need to update the network for the actions |
| 330 | # which were actually taken. |
| 331 | # We can accomplish this by setting the target to be equal to |
| 332 | # the prediction for all values. |
| 333 | # Then, only change the targets for the actions taken. |
| 334 | # Q(s,a) |
| 335 | target_full = predict(self.model, states) |
| 336 | target_full[np.arange(batch_size), actions] = target |
| 337 | |
| 338 | # Run one training step |
| 339 | train_one_step(self.model, self.criterion, self.optimizer, states, target_full) |
| 340 | |
| 341 | if self.epsilon > self.epsilon_min: |