-
Notifications
You must be signed in to change notification settings - Fork 0
/
algorithm.py
63 lines (46 loc) · 2.13 KB
/
algorithm.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
from torch import nn
from torch.optim import Adam
import torch
import numpy as np
learning_rate = 1e-2
GAMMA = 0.95
class REINFORCE:
def __init__(self, input_dim, hidden_dim, action_dim, max_len):
self.model = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, action_dim),
nn.Softmax(dim=-1))
self.action_dim = action_dim
self.optimizer = Adam(self.model.parameters(), learning_rate)
self.gammas = torch.tensor([np.power(GAMMA, i) for i in range(max_len)])
def update(self, trajectories):
total_sum = torch.Tensor(len(trajectories))
for j,traj in enumerate(trajectories):
# Take cumulative rewards with discount factor
rewards = torch.tensor([moment[3] for moment in reversed(traj)])
gammas_reversed = torch.flip(self.gammas[:len(traj)], [0])
cum_reward = torch.cumsum(rewards * gammas_reversed, dim=0)
cum_reward = torch.flip(cum_reward, [0])
cum_reward = cum_reward / self.gammas[:len(traj)]
# Get probabilities
output = self.model(torch.tensor([moment[0] for moment in traj]).float())
# Take logarithm of probabilities
prob_ac = torch.Tensor(len(traj))
for i,m in enumerate(traj):
prob_ac[i] = output[i, m[1]]
prob_log = torch.log(prob_ac)
# Get loss
total_sum[j] = torch.sum(prob_log * cum_reward)
# Take a gradient step
loss = -torch.mean(total_sum)
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
def action(self, state):
state = torch.tensor(state).float()
distribution = self.model(state)
ac = np.random.choice(range(self.action_dim), p=distribution.detach().numpy())
return ac