--- a +++ b/SAC/model.py @@ -0,0 +1,205 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.distributions import Normal +from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence + +LOG_SIG_MAX = 2 +LOG_SIG_MIN = -20 +epsilon = 1e-6 + +# Initialize Policy weights +def weights_init_(m): + if isinstance(m, nn.Linear): + torch.nn.init.xavier_uniform_(m.weight, gain=1) + torch.nn.init.constant_(m.bias, 0) + +class Actor(nn.Module): + def __init__(self, num_inputs, num_actions, hidden_dim, model, action_space=None): + super(Actor, self).__init__() + + self.linear1 = nn.Linear(num_inputs, hidden_dim) + + if model == "rnn": + self.rnn = nn.RNN(hidden_dim, hidden_dim, batch_first=True) + elif model == "gru": + self.rnn = nn.GRU(hidden_dim, hidden_dim, batch_first=True) + else: + raise NotImplementedError + + self.mean_linear = nn.Linear(hidden_dim, num_actions) + self.log_std_linear = nn.Linear(hidden_dim, num_actions) + + self.apply(weights_init_) + + # action rescaling + # Pass none action space and adjust the action scale and bias manually + if action_space is None: + self.action_scale = torch.tensor(0.5) + self.action_bias = torch.tensor(0.5) + else: + self.action_scale = torch.FloatTensor( + (action_space.high - action_space.low) / 2.) + self.action_bias = torch.FloatTensor( + (action_space.high + action_space.low) / 2.) + + def forward(self, state, h_prev, sampling, len_seq= None): + + x = F.tanh(self.linear1(state)) + + if sampling == False: + assert len_seq!=None, "Proved the len_seq" + x = pack_padded_sequence(x, len_seq, batch_first= True, enforce_sorted= False) + + #Tap RNN input for fixedpoint analysis + rnn_in = x + + x, (h_current) = self.rnn(x, (h_prev)) + + if sampling == False: + x, _ = pad_packed_sequence(x, batch_first= True) + + if sampling == True: + x = x.squeeze(1) + + # x = F.relu(x) + mean = self.mean_linear(x) + log_std = self.log_std_linear(x) + log_std = torch.clamp(log_std, min=LOG_SIG_MIN, max=LOG_SIG_MAX) + + return mean, log_std, h_current, x, rnn_in + + def sample(self, state, h_prev, sampling, len_seq=None): + + mean, log_std, h_current, x, rnn_in = self.forward(state, h_prev, sampling, len_seq) + #if sampling == False; then mask the mean and log_std using len_seq + if sampling == False: + assert mean.size()[1] == log_std.size()[1], "There is a mismatch between and mean and sigma Sl_max" + sl_max = mean.size()[1] + with torch.no_grad(): + for seq_idx, k in enumerate(len_seq): + for j in range(1, sl_max + 1): + if j <= k: + if seq_idx == 0 and j == 1: + mask_seq = torch.tensor([True], dtype=bool) + else: + mask_seq = torch.cat((mask_seq, torch.tensor([True])), dim=0) + else: + mask_seq = torch.cat((mask_seq, torch.tensor([False])), dim=0) + #The mask has been created, Now filter the mean and sigma using this mask + mean = mean.reshape(-1, mean.size()[-1])[mask_seq] + log_std = log_std.reshape(-1, log_std.size()[-1])[mask_seq] + + if sampling == True: + mask_seq = [] #If sampling is True return a dummy mask seq + + std = log_std.exp() + + # white noise + normal = Normal(mean, std) + noise = normal.rsample() + + # reparameterization trick + y_t = torch.tanh(noise) + action = y_t * self.action_scale + self.action_bias + log_prob = normal.log_prob(noise) + + # Enforce the action_bounds + log_prob -= torch.log(self.action_scale * (1 - y_t.pow(2)) + epsilon) + log_prob = log_prob.sum(1, keepdim=True) + mean = torch.tanh(mean) * self.action_scale + self.action_bias + + return action, log_prob, mean, h_current, mask_seq, x, rnn_in + + def forward_for_simple_dynamics(self, state, h_prev, sampling, len_seq= None): + + x = F.tanh(self.linear1(state)) + + #Tap the output of the first linear layer + x_l1 = x + + if sampling == False: + assert len_seq!=None, "Proved the len_seq" + x = pack_padded_sequence(x, len_seq, batch_first= True, enforce_sorted= False) + + x, _ = self.rnn(x, (h_prev)) + + if sampling == False: + x, _ = pad_packed_sequence(x, batch_first= True) + + # x = F.relu(x) + + return x, x_l1 + + + def forward_lstm(self, state, h_prev, sampling, len_seq= None): + + x = F.tanh(self.linear1(state)) + + if sampling == False: + assert len_seq!=None, "Proved the len_seq" + + x = pack_padded_sequence(x, len_seq, batch_first= True, enforce_sorted= False) + + x, (h_current) = self.rnn(x, (h_prev)) + + if sampling == False: + x, len_x_seq = pad_packed_sequence(x, batch_first= True) + + if sampling == True: + x = x.squeeze(1) + + return x + + def forward_for_neural_pert(self, state, h_prev, neural_pert= None): + + x = F.tanh(self.linear1(state)) + + #Tap RNN input for fixedpoint analysis + rnn_in = x + + x, (h_current) = self.rnn(x, (h_prev)) + + #Add the neural perturbation to the RNN output + x = x+neural_pert + + x = x.squeeze(1) + + mean = self.mean_linear(x) + log_std = self.log_std_linear(x) + log_std = torch.clamp(log_std, min=LOG_SIG_MIN, max=LOG_SIG_MAX) + + + action = torch.tanh(mean) * self.action_scale + self.action_bias + + return action.detach().cpu().numpy()[0], h_current.detach(), x.detach().cpu().numpy(), rnn_in.detach().cpu().numpy() + +class Critic(nn.Module): + def __init__(self, num_inputs, num_actions, hidden_dim): + super(Critic, self).__init__() + + # Q1 architecture + self.linear1 = nn.Linear(num_inputs + num_actions, hidden_dim) + self.linear2 = nn.Linear(hidden_dim, hidden_dim) + self.linear3 = nn.Linear(hidden_dim, 1) + + # Q2 architecture + self.linear4 = nn.Linear(num_inputs + num_actions, hidden_dim) + self.linear5 = nn.Linear(hidden_dim, hidden_dim) + self.linear6 = nn.Linear(hidden_dim, 1) + + self.apply(weights_init_) + + def forward(self, state, action): + + xu = torch.cat([state, action], 1) + + x1 = F.relu(self.linear1(xu)) + x1 = F.relu(self.linear2(x1)) + x1 = self.linear3(x1) + + x2 = F.relu(self.linear4(xu)) + x2 = F.relu(self.linear5(x2)) + x2 = self.linear6(x2) + + return x1, x2 \ No newline at end of file