import numpy as np
import torch
from SAC.sac import SAC_Agent
from SAC.replay_memory import PolicyReplayMemory
from SAC.RL_Framework_Mujoco import Muscle_Env
from SAC import sensory_feedback_specs, kinematics_preprocessing_specs, perturbation_specs
import pickle
import os
from numpy.core.records import fromarrays
from scipy.io import savemat
#Set the current working directory
os.chdir(os.getcwd())
class Simulate():
def __init__(self, env:Muscle_Env, args):
"""Train a soft actor critic agent to control a musculoskeletal model to follow a kinematic trajectory.
Parameters
----------
env: str
specify which environment (model) to train on
model: str
specify whether to use an rnn or a gated rnn (gru)
optimizer_spec: float
gamma discount parameter in sac algorithm
tau: float
tau parameter for soft updates of the critic and critic target networks
lr: float
learning rate of the critic and actor networks
alpha: float
entropy term used in the sac policy update
automatic entropy tuning: bool
whether to use automatic entropy tuning during training
seed: int
seed for the environment and networks
policy_batch_size: int
the number of episodes used in a batch during training
hidden_size: int
number of hidden neurons in the actor and critic
policy_replay_size: int
number of episodes to store in the replay
multi_policy_loss: bool
use additional policy losses during updates (reduce norm of weights)
ONLY USE WITH RNN, NOT IMPLEMENTED WITH GATING
batch_iters: int
How many experience replay rounds (not steps!) to perform between
each update
cuda: bool
use cuda gpu
visualize: bool
visualize model
model_save_name: str
specify the name of the model to save
episodes: int
total number of episodes to run
save_iter: int
number of iterations before saving the model
checkpoint_path: str
specify path to save and load model
muscle_path: str
path for the musculoskeletal model
muscle_params_path: str
path for musculoskeletal model parameters
"""
### TRAINING VARIABLES ###
self.episodes = args.total_episodes
self.hidden_size = args.hidden_size
self.policy_batch_size = args.policy_batch_size
self.visualize = args.visualize
self.root_dir = args.root_dir
self.checkpoint_file = args.checkpoint_file
self.checkpoint_folder = args.checkpoint_folder
self.statistics_folder = args.statistics_folder
self.batch_iters = args.batch_iters
self.save_iter = args.save_iter
self.mode_to_sim = args.mode
self.condition_selection_strategy = args.condition_selection_strategy
self.load_saved_nets_for_training = args.load_saved_nets_for_training
self.verbose_training = args.verbose_training
### ENSURE SAVING FILES ARE ACCURATE ###
assert isinstance(self.root_dir, str)
assert isinstance(self.checkpoint_folder, str)
assert isinstance(self.checkpoint_file, str)
### SEED ###
torch.manual_seed(args.seed)
np.random.seed(args.seed)
### LOAD CUSTOM GYM ENVIRONMENT ###
if self.mode_to_sim in ["musculo_properties"]:
self.env = env(args.musculoskeletal_model_path[:-len('musculoskeletal_model.xml')] + 'musculo_targets_pert.xml', 1, args)
else:
self.env = env(args.musculoskeletal_model_path[:-len('musculoskeletal_model.xml')] + 'musculo_targets.xml', 1, args)
self.observation_shape = self.env.observation_space.shape[0]+len(self.env.sfs_visual_velocity)*3+1
### SAC AGENT ###
self.agent = SAC_Agent(self.observation_shape,
self.env.action_space,
args.hidden_size,
args.lr,
args.gamma,
args.tau,
args.alpha,
args.automatic_entropy_tuning,
args.model,
args.multi_policy_loss,
args.alpha_usim,
args.beta_usim,
args.gamma_usim,
args.zeta_nusim,
args.cuda)
### REPLAY MEMORY ###
self.policy_memory = PolicyReplayMemory(args.policy_replay_size, args.seed)
def test(self, save_name):
""" Use a saved model to generate kinematic trajectory
saves these values:
--------
episode_reward: int
- Total reward of the trained model
kinematics: list
- Kinematics of the model during testing
hidden_activity: list
- list containing the rnn activity during testing
"""
#Update the environment kinematics to both the training and testing conditions
self.env.update_kinematics_for_test()
### LOAD SAVED MODEL ###
self.load_saved_nets_from_checkpoint(load_best = True)
#Set the recurrent connections to zero if the mode is SFE
if self.mode_to_sim in ["SFE"] and "recurrent_connections" in perturbation_specs.sf_elim:
self.agent.actor.rnn.weight_hh_l0 = torch.nn.Parameter(self.agent.actor.rnn.weight_hh_l0 * 0)
### TESTING DATA ###
Test_Data = {
"emg": {},
"rnn_activity": {},
"rnn_input": {},
"rnn_input_fp": {},
"kinematics_mbodies": {},
"kinematics_mtargets": {}
}
###Data jPCA
activity_jpca = []
times_jpca = []
n_fixedsteps_jpca = []
condition_tpoints_jpca = []
#Save the following after testing: EMG, RNN_activity, RNN_input, kin_musculo_bodies, kin_musculo_targets
emg = {}
rnn_activity = {}
rnn_input = {}
kin_mb = {}
kin_mt = {}
rnn_input_fp = {}
for i_cond_sim in range(self.env.n_exp_conds):
### TRACKING VARIABLES ###
episode_reward = 0
episode_steps = 0
emg_cond = []
kin_mb_cond = []
kin_mt_cond = []
rnn_activity_cond = []
rnn_input_cond = []
rnn_input_fp_cond = []
done = False
### GET INITAL STATE + RESET MODEL BY POSE
cond_to_select = i_cond_sim % self.env.n_exp_conds
state = self.env.reset(cond_to_select)
if self.mode_to_sim in ["SFE"] and "task_scalar" in perturbation_specs.sf_elim:
state = [*state, 0]
else:
state = [*state, self.env.condition_scalar]
# Num_layers specified in the policy model
h_prev = torch.zeros(size=(1, 1, self.hidden_size))
### STEPS PER EPISODE ###
for timestep in range(self.env.timestep_limit):
### SELECT ACTION ###
with torch.no_grad():
if self.mode_to_sim in ["neural_pert"]:
state = torch.FloatTensor(state).to(self.agent.device).unsqueeze(0).unsqueeze(0)
h_prev = h_prev.to(self.agent.device)
neural_pert = perturbation_specs.neural_pert[timestep % perturbation_specs.neural_pert.shape[0], :]
neural_pert = torch.FloatTensor(neural_pert).to(self.agent.device).unsqueeze(0).unsqueeze(0)
action, h_current, rnn_act, rnn_in = self.agent.actor.forward_for_neural_pert(state, h_prev, neural_pert)
elif self.mode_to_sim in ["SFE"] and "recurrent_connections" in perturbation_specs.sf_elim:
h_prev = h_prev*0
action, h_current, rnn_act, rnn_in = self.agent.select_action(state, h_prev, evaluate=True)
else:
action, h_current, rnn_act, rnn_in = self.agent.select_action(state, h_prev, evaluate=True)
emg_cond.append(action) #[n_muscles, ]
rnn_activity_cond.append(rnn_act[0, :]) # [1, n_hidden_units] --> [n_hidden_units,]
rnn_input_cond.append(state) #[n_inputs, ]
rnn_input_fp_cond.append(rnn_in[0, 0, :]) #[1, 1, n_hidden_units] --> [n_hidden_units, ]
### TRACKING REWARD + EXPERIENCE TUPLE###
next_state, reward, done, _ = self.env.step(action)
if self.mode_to_sim in ["SFE"] and "task_scalar" in perturbation_specs.sf_elim:
next_state = [*next_state, 0]
else:
next_state = [*next_state, self.env.condition_scalar]
episode_reward += reward
#now append the kinematics of the musculo body and the corresponding target
kin_mb_t = []
kin_mt_t = []
for musculo_body in kinematics_preprocessing_specs.musculo_tracking:
kin_mb_t.append(self.env.sim.data.get_body_xpos(musculo_body[0]).copy()) #[3, ]
kin_mt_t.append(self.env.sim.data.get_body_xpos(musculo_body[1]).copy()) #[3, ]
kin_mb_cond.append(kin_mb_t) # kin_mb_t : [n_targets, 3]
kin_mt_cond.append(kin_mt_t) # kin_mt_t : [n_targets, 3]
### VISUALIZE MODEL ###
if self.visualize == True:
self.env.render()
state = next_state
h_prev = h_current
#Append the testing data
emg[i_cond_sim] = np.array(emg_cond) # [timepoints, muscles]
rnn_activity[i_cond_sim] = np.array(rnn_activity_cond) # [timepoints, n_hidden_units]
rnn_input[i_cond_sim] = np.array(rnn_input_cond) #[timepoints, n_inputs]
rnn_input_fp[i_cond_sim] = np.array(rnn_input_fp_cond) #[timepoints, n_hidden_units]
kin_mb[i_cond_sim] = np.array(kin_mb_cond).transpose(1, 0, 2) # kin_mb_cond: [timepoints, n_targets, 3] --> [n_targets, timepoints, 3]
kin_mt[i_cond_sim] = np.array(kin_mt_cond).transpose(1, 0, 2) # kin_mt_cond: [timepoints, n_targets, 3] --> [n_targets, timepoints, 3]
##Append the jpca data
activity_jpca.append(dict(A = rnn_activity_cond))
times_jpca.append(dict(times = np.arange(timestep))) #the timestep is assumed to be 1ms
condition_tpoints_jpca.append(self.env.kin_to_sim[self.env.current_cond_to_sim].shape[-1])
n_fixedsteps_jpca.append(self.env.n_fixedsteps)
### SAVE TESTING STATS ###
Test_Data["emg"] = emg
Test_Data["rnn_activity"] = rnn_activity
Test_Data["rnn_input"] = rnn_input
Test_Data["rnn_input_fp"] = rnn_input_fp
Test_Data["kinematics_mbodies"] = kin_mb
Test_Data["kinematics_mtargets"] = kin_mt
### Save the jPCA data
Data_jpca = fromarrays([activity_jpca, times_jpca], names=['A', 'times'])
#save test data
with open(save_name + '/test_data.pkl', 'wb') as f:
pickle.dump(Test_Data, f)
#save jpca data
savemat(save_name + '/Data_jpca.mat', {'Data' : Data_jpca})
savemat(save_name + '/n_fixedsteps_jpca.mat', {'n_fsteps' : n_fixedsteps_jpca})
savemat(save_name + '/condition_tpoints_jpca.mat', {'cond_tpoints': condition_tpoints_jpca})
def train(self):
""" Train an RNN based SAC agent to follow kinematic trajectory
"""
#Load the saved networks from the last training
if self.load_saved_nets_for_training:
self.load_saved_nets_from_checkpoint(load_best= False)
### TRAINING DATA DICTIONARY ###
Statistics = {
"rewards": [],
"steps": [],
"policy_loss": [],
"critic_loss": []
}
highest_reward = -float("inf") # used for storing highest reward throughout training
#Average reward across conditions initialization
cond_train_count= np.ones((self.env.n_exp_conds,))
cond_avg_reward = np.zeros((self.env.n_exp_conds,))
cond_cum_reward = np.zeros((self.env.n_exp_conds,))
cond_cum_count = np.zeros((self.env.n_exp_conds,))
### BEGIN TRAINING ###
for episode in range(self.episodes):
### Gather Episode Data Variables ###
episode_reward = 0 # reward for single episode
episode_steps = 0 # steps completed for single episode
policy_loss_tracker = [] # stores policy loss throughout episode
critic1_loss_tracker = [] # stores critic loss throughout episode
done = False # determines if episode is terminated
### GET INITAL STATE + RESET MODEL BY POSE
if self.condition_selection_strategy != "reward":
cond_to_select = episode % self.env.n_exp_conds
state = self.env.reset(cond_to_select)
else:
cond_indx = np.nonzero(cond_train_count>0)[0][0]
cond_train_count[cond_indx] = cond_train_count[cond_indx] - 1
state = self.env.reset(cond_indx)
#Append the high-level task scalar signal
state = [*state, self.env.condition_scalar]
ep_trajectory = [] # used to store (s_t, a_t, r_t, s_t+1) tuple for replay storage
h_prev = torch.zeros(size=(1, 1, self.hidden_size)) # num_layers specified in the policy model
### LOOP THROUGH EPISODE TIMESTEPS ###
while not(done):
### SELECT ACTION ###
with torch.no_grad():
action, h_current, _, _ = self.agent.select_action(state, h_prev, evaluate=False)
#Now query the neural activity idx from the simulator
na_idx= self.env.coord_idx
### UPDATE MODEL PARAMETERS ###
if len(self.policy_memory.buffer) > self.policy_batch_size:
for _ in range(self.batch_iters):
critic_1_loss, critic_2_loss, policy_loss = self.agent.update_parameters(self.policy_memory, self.policy_batch_size)
### STORE LOSSES ###
policy_loss_tracker.append(policy_loss)
critic1_loss_tracker.append(critic_1_loss)
### SIMULATION ###
reward = 0
for _ in range(self.env.frame_repeat):
next_state, inter_reward, done, _ = self.env.step(action)
next_state = [*next_state, self.env.condition_scalar]
reward += inter_reward
episode_steps += 1
### EARLY TERMINATION OF EPISODE ###
if done:
break
episode_reward += reward
### VISUALIZE MODEL ###
if self.visualize == True:
self.env.render()
mask = 1 if episode_steps == self.env._max_episode_steps else float(not done) # ensure mask is not 0 if episode ends
### STORE CURRENT TIMESTEP TUPLE ###
if not self.env.nusim_data_exists:
self.env.neural_activity[na_idx] = 0
ep_trajectory.append((state,
action,
reward,
next_state,
mask,
h_current.squeeze(0).cpu().numpy(),
self.env.neural_activity[na_idx],
np.array([na_idx])))
### MOVE TO NEXT STATE ###
state = next_state
h_prev = h_current
### PUSH TO REPLAY ###
self.policy_memory.push(ep_trajectory)
if self.condition_selection_strategy == "reward":
cond_cum_reward[cond_indx] = cond_cum_reward[cond_indx] + episode_reward
cond_cum_count[cond_indx] = cond_cum_count[cond_indx] + 1
#Check if there are all zeros in the cond_train_count array
if np.all((cond_train_count == 0)):
cond_avg_reward = cond_cum_reward / cond_cum_count
cond_train_count = np.ceil((np.max(cond_avg_reward)*np.ones((self.env.n_exp_conds,)))/cond_avg_reward)
### TRACKING ###
Statistics["rewards"].append(episode_reward)
Statistics["steps"].append(episode_steps)
Statistics["policy_loss"].append(np.mean(np.array(policy_loss_tracker)))
Statistics["critic_loss"].append(np.mean(np.array(critic1_loss_tracker)))
### SAVE DATA TO FILE (in root project folder) ###
if len(self.statistics_folder) != 0:
np.save(self.statistics_folder + f'/stats_rewards.npy', Statistics['rewards'])
np.save(self.statistics_folder + f'/stats_steps.npy', Statistics['steps'])
np.save(self.statistics_folder + f'/stats_policy_loss.npy', Statistics['policy_loss'])
np.save(self.statistics_folder + f'/stats_critic_loss.npy', Statistics['critic_loss'])
### SAVING STATE DICT OF TRAINING ###
if len(self.checkpoint_folder) != 0 and len(self.checkpoint_file) != 0:
if episode % self.save_iter == 0 and len(self.policy_memory.buffer) > self.policy_batch_size:
#Save the state dicts
torch.save({
'iteration': episode,
'agent_state_dict': self.agent.actor.state_dict(),
'critic_state_dict': self.agent.critic.state_dict(),
'critic_target_state_dict': self.agent.critic_target.state_dict(),
'agent_optimizer_state_dict': self.agent.actor_optim.state_dict(),
'critic_optimizer_state_dict': self.agent.critic_optim.state_dict(),
}, self.checkpoint_folder + f'/{self.checkpoint_file}.pth')
#Save the pickled model for fixedpoint finder analysis
torch.save(self.agent.actor.rnn, self.checkpoint_folder + f'/actor_rnn_fpf.pth')
if episode_reward > highest_reward:
torch.save({
'iteration': episode,
'agent_state_dict': self.agent.actor.state_dict(),
'critic_state_dict': self.agent.critic.state_dict(),
'critic_target_state_dict': self.agent.critic_target.state_dict(),
'agent_optimizer_state_dict': self.agent.actor_optim.state_dict(),
'critic_optimizer_state_dict': self.agent.critic_optim.state_dict(),
}, self.checkpoint_folder + f'/{self.checkpoint_file}_best.pth')
#Save the pickled model for fixedpoint finder analysis
torch.save(self.agent.actor.rnn, self.checkpoint_folder + f'/actor_rnn_best_fpf.pth')
if episode_reward > highest_reward:
highest_reward = episode_reward
### PRINT TRAINING OUTPUT ###
if self.verbose_training:
print('-----------------------------------')
print('highest reward: {} | reward: {} | timesteps completed: {}'.format(highest_reward, episode_reward, episode_steps))
print('-----------------------------------\n')
def load_saved_nets_from_checkpoint(self, load_best: bool):
#Load the saved networks from the checkpoint file
#Saved networks include policy, critic, critic_target, policy_optimizer and critic_optimizer
if not load_best:
#Load the policy network
self.agent.actor.load_state_dict(torch.load(self.checkpoint_folder + f'/{self.checkpoint_file}.pth')['agent_state_dict'])
#Load the critic network
self.agent.critic.load_state_dict(torch.load(self.checkpoint_folder + f'/{self.checkpoint_file}.pth')['critic_state_dict'])
#Load the critic target network
self.agent.critic_target.load_state_dict(torch.load(self.checkpoint_folder + f'/{self.checkpoint_file}.pth')['critic_target_state_dict'])
#Load the policy optimizer
self.agent.actor_optim.load_state_dict(torch.load(self.checkpoint_folder + f'/{self.checkpoint_file}.pth')['agent_optimizer_state_dict'])
#Load the critic optimizer
self.agent.critic_optim.load_state_dict(torch.load(self.checkpoint_folder + f'/{self.checkpoint_file}.pth')['critic_optimizer_state_dict'])
else:
#Load the policy network
self.agent.actor.load_state_dict(torch.load(self.checkpoint_folder + f'/{self.checkpoint_file}_best.pth')['agent_state_dict'])
#Load the critic network
self.agent.critic.load_state_dict(torch.load(self.checkpoint_folder + f'/{self.checkpoint_file}_best.pth')['critic_state_dict'])
#Load the critic target network
self.agent.critic_target.load_state_dict(torch.load(self.checkpoint_folder + f'/{self.checkpoint_file}_best.pth')['critic_target_state_dict'])
#Load the policy optimizer
self.agent.actor_optim.load_state_dict(torch.load(self.checkpoint_folder + f'/{self.checkpoint_file}_best.pth')['agent_optimizer_state_dict'])
#Load the critic optimizer
self.agent.critic_optim.load_state_dict(torch.load(self.checkpoint_folder + f'/{self.checkpoint_file}_best.pth')['critic_optimizer_state_dict'])
Datasets
Models
