import os
import torch
import torch.nn.functional as F
from torch.optim import Adam
from .utils import soft_update, hard_update
from .model import Actor, Critic
from torch.nn.utils.rnn import pad_sequence
import numpy as np
from .replay_memory import PolicyReplayMemory
import ipdb
class SAC_Agent():
def __init__(self,
num_inputs: int,
action_space: int,
hidden_size: int,
lr: float,
gamma: float,
tau: float,
alpha: float,
automatic_entropy_tuning: bool,
model: str,
multi_policy_loss: bool,
alpha_usim:float,
beta_usim:float,
gamma_usim:float,
zeta_nusim:float,
cuda: bool):
if cuda:
self.device = torch.device("cuda")
else:
self.device = torch.device("cpu")
#Save the regularization parameters
self.alpha_usim = alpha_usim
self.beta_usim = beta_usim
self.gamma_usim = gamma_usim
self.zeta_nusim = zeta_nusim
### SET CRITIC NETWORKS ###
self.critic = Critic(num_inputs, action_space.shape[0], hidden_size).to(self.device)
self.critic_target = Critic(num_inputs, action_space.shape[0], hidden_size).to(self.device)
self.critic_optim = Adam(self.critic.parameters(), lr=lr)
hard_update(self.critic_target, self.critic)
### SET ACTOR NETWORK ###
self.actor = Actor(num_inputs, action_space.shape[0], hidden_size, model, action_space=None).to(self.device)
self.actor_optim = Adam(self.actor.parameters(), lr=lr)
### SET TRAINING VARIABLES ###
self.model = model
self.multi_policy_loss = multi_policy_loss
self.gamma = gamma
self.tau = tau
self.alpha = alpha
self.hidden_size= hidden_size
self.automatic_entropy_tuning = automatic_entropy_tuning
# Target Entropy = −dim(A) (e.g. , -6 for HalfCheetah-v2) as given in the paper
if automatic_entropy_tuning:
self.target_entropy = -torch.prod(torch.Tensor(action_space.shape).to(self.device)).item()
self.log_alpha = torch.zeros(1, requires_grad=True, device=self.device)
self.alpha_optim = Adam([self.log_alpha], lr=lr)
#This loss encourages the simple low-dimensional dynamics in the RNN activity
def _policy_loss_2(self, policy_state_batch, h0, len_seq, mask_seq):
# Sample the hidden weights of the RNN
J_rnn_w = self.actor.rnn.weight_hh_l0 #These weights would be of the size (hidden_dim, hidden_dim)
#Sample the output of the RNN for the policy_state_batch
rnn_out_r, _ = self.actor.forward_for_simple_dynamics(policy_state_batch, h0, sampling=False, len_seq=len_seq)
rnn_out_r = rnn_out_r.reshape(-1, rnn_out_r.size()[-1])[mask_seq]
#Reshape the policy hidden weights vector
J_rnn_w = J_rnn_w.unsqueeze(0).repeat(rnn_out_r.size()[0], 1, 1)
rnn_out_r = 1 - torch.pow(rnn_out_r, 2)
R_j = torch.mul(J_rnn_w, rnn_out_r.unsqueeze(-1))
policy_loss_2 = torch.norm(R_j)**2
return policy_loss_2
#This loss encourages the minimization of the firing rates for the linear and the RNN layer.
def _policy_loss_3(self, policy_state_batch, h0, len_seq, mask_seq):
#Find the loss encouraging the minimization of the firing rates for the linear and the RNN layer
#Sample the output of the RNN for the policy_state_batch
rnn_out_r, linear_out = self.actor.forward_for_simple_dynamics(policy_state_batch, h0, sampling=False, len_seq=len_seq)
rnn_out_r = rnn_out_r.reshape(-1, rnn_out_r.size()[-1])[mask_seq]
linear_out = linear_out.reshape(-1, linear_out.size()[-1])[mask_seq]
policy_loss_3 = torch.norm(rnn_out_r)**2 + torch.norm(linear_out)**2
return policy_loss_3
#This loss encourages the minimization of the input and output weights of the RNN and the layers downstream/
#upstream of the RNN.
def _policy_loss_4(self):
#Find the loss encouraging the minimization of the input and output weights of the RNN and the layers downstream
#and upstream of the RNN
#Sample the input weights of the RNN
J_rnn_i = self.actor.rnn.weight_ih_l0
J_in1 = self.actor.linear1.weight
#Sample the output weights
J_out1 = self.actor.mean_linear.weight
J_out2 = self.actor.log_std_linear.weight
policy_loss_4 = torch.norm(J_in1)**2 + torch.norm(J_rnn_i)**2 + torch.norm(J_out1)**2 + torch.norm(J_out2)**2
return policy_loss_4
#Define a loss function that constraints a subset of the RNN nodes to the experimental neural data
def _policy_loss_exp_neural_constrain(self, policy_state_batch, h0, len_seq, neural_activity_batch, na_idx_batch, mask_seq):
#Find the loss for neural activity constrainting
lstm_out = self.actor.forward_lstm(policy_state_batch, h0, sampling= False, len_seq= len_seq)
lstm_out = lstm_out.reshape(-1, lstm_out.size()[-1])[mask_seq]
lstm_activity = lstm_out[:, 0:neural_activity_batch.shape[-1]]
#Now filter the neural activity batch and lstm activity batch using the na_idx batch
with torch.no_grad():
na_idx_batch = na_idx_batch.squeeze(-1) > 0
lstm_activity = lstm_activity[na_idx_batch]
neural_activity_batch = neural_activity_batch[na_idx_batch]
policy_loss_exp_c = F.mse_loss(lstm_activity, neural_activity_batch)
return policy_loss_exp_c
def select_action(self, state: np.ndarray, h_prev: torch.Tensor, evaluate=False) -> (np.ndarray, torch.Tensor, np.ndarray):
state = torch.FloatTensor(state).to(self.device).unsqueeze(0).unsqueeze(0)
h_prev = h_prev.to(self.device)
### IF TRAINING ###
if evaluate == False:
# get action sampled from gaussian
action, _, _, h_current, _, rnn_out, rnn_in = self.actor.sample(state, h_prev, sampling=True, len_seq=None)
### IF TESTING ###
else:
# get the action without noise
_, _, action, h_current, _, rnn_out, rnn_in = self.actor.sample(state, h_prev, sampling=True, len_seq=None)
return action.detach().cpu().numpy()[0], h_current.detach(), rnn_out.detach().cpu().numpy(), rnn_in.detach().cpu().numpy()
def update_parameters(self, policy_memory: PolicyReplayMemory, policy_batch_size: int) -> (int, int, int):
### SAMPLE FROM REPLAY ###
state_batch, action_batch, reward_batch, next_state_batch, mask_batch, h_batch, policy_state_batch, neural_activity_batch, na_idx_batch = policy_memory.sample(batch_size=policy_batch_size)
### CONVERT DATA TO TENSOR ###
state_batch = torch.FloatTensor(state_batch).to(self.device)
next_state_batch = torch.FloatTensor(next_state_batch).to(self.device)
action_batch = torch.FloatTensor(action_batch).to(self.device)
reward_batch = torch.FloatTensor(reward_batch).to(self.device).unsqueeze(1)
mask_batch = torch.FloatTensor(mask_batch).to(self.device).unsqueeze(1)
h_batch = torch.FloatTensor(h_batch).to(self.device).permute(1, 0, 2)
neural_activity_batch = torch.FloatTensor(neural_activity_batch).to(self.device)
na_idx_batch = torch.FloatTensor(na_idx_batch).to(self.device)
h0 = torch.zeros(size=(1, next_state_batch.shape[0], self.hidden_size)).to(self.device)
### SAMPLE NEXT Q VALUE FOR CRITIC LOSS ###
with torch.no_grad():
next_state_action, next_state_log_pi, _, _, _, _, _ = self.actor.sample(next_state_batch.unsqueeze(1), h_batch, sampling=True)
qf1_next_target, qf2_next_target = self.critic_target(next_state_batch, next_state_action)
min_qf_next_target = torch.min(qf1_next_target, qf2_next_target) - self.alpha * next_state_log_pi
next_q_value = reward_batch + mask_batch * self.gamma * (min_qf_next_target)
### CALCULATE CRITIC LOSS ###
qf1, qf2 = self.critic(state_batch, action_batch) # Two Q-functions to mitigate positive bias in the policy improvement step
qf1_loss = F.mse_loss(qf1, next_q_value) # JQ = 𝔼(st,at)~D[0.5(Q1(st,at) - r(st,at) - γ(𝔼st+1~p[V(st+1)]))^2]
qf2_loss = F.mse_loss(qf2, next_q_value) # JQ = 𝔼(st,at)~D[0.5(Q1(st,at) - r(st,at) - γ(𝔼st+1~p[V(st+1)]))^2]
qf_loss = qf1_loss + qf2_loss
### TAKE GRAIDENT STEP ###
self.critic_optim.zero_grad()
qf_loss.backward()
self.critic_optim.step()
### SAMPLE FROM ACTOR NETWORK ###
h0 = torch.zeros(size=(1, len(policy_state_batch), self.hidden_size)).to(self.device)
len_seq = list(map(len, policy_state_batch))
policy_state_batch = torch.FloatTensor(pad_sequence(policy_state_batch, batch_first= True)).to(self.device)
pi_action_bat, log_prob_bat, _, _, mask_seq, _, _ = self.actor.sample(policy_state_batch, h0, sampling=False, len_seq=len_seq)
### MASK POLICY STATE BATCH ###
policy_state_batch_pi = policy_state_batch.reshape(-1, policy_state_batch.size()[-1])[mask_seq]
### GET VALUE OF CURRENT STATE AND ACTION PAIRS ###
qf1_pi, qf2_pi = self.critic(policy_state_batch_pi, pi_action_bat)
min_qf_pi = torch.min(qf1_pi, qf2_pi)
### CALCULATE POLICY LOSS ###
task_loss = ((self.alpha * log_prob_bat) - min_qf_pi).mean() # Jπ = 𝔼st∼D,εt∼N[α * logπ(f(εt;st)|st) − Q(st,f(εt;st))]
policy_loss = task_loss
############################
# ADDITIONAL POLICY LOSSES #
############################
if self.multi_policy_loss:
loss_simple_dynamics = self._policy_loss_2(policy_state_batch, h0, len_seq, mask_seq)
loss_activations_min = self._policy_loss_3(policy_state_batch, h0, len_seq, mask_seq)
loss_weights_min = self._policy_loss_4()
loss_exp_constrain = self._policy_loss_exp_neural_constrain(policy_state_batch, h0, len_seq, neural_activity_batch, na_idx_batch, mask_seq)
### CALCULATE FINAL POLICY LOSS ###
#To implement nuSim training use a weighting of 1e+04 with loss_exp_constrain
policy_loss += (self.alpha_usim*(loss_simple_dynamics)) \
+ (self.beta_usim*(loss_activations_min)) \
+ (self.gamma_usim*(loss_weights_min)) \
+ (self.zeta_nusim*(loss_exp_constrain))
### TAKE GRADIENT STEP ###
self.actor_optim.zero_grad()
policy_loss.backward()
self.actor_optim.step()
### AUTOMATIC ENTROPY TUNING ###
log_pi = log_prob_bat
if self.automatic_entropy_tuning:
alpha_loss = -(self.log_alpha * (log_pi + self.target_entropy).detach()).mean()
self.alpha_optim.zero_grad()
alpha_loss.backward()
self.alpha_optim.step()
self.alpha = self.log_alpha.exp()
### SOFT UPDATE OF CRITIC TARGET ###
soft_update(self.critic_target, self.critic, self.tau)
return qf1_loss.item(), qf2_loss.item(), policy_loss.item()
Datasets
Models
