class SAC(object):
def __init__(self, num_inputs, action_space, args):
self.gamma = args.gamma
self.tau = args.tau
self.alpha = args.alpha
self.policy_type = args.policy
self.target_update_interval = args.target_update_interval
self.automatic_entropy_tuning = args.automatic_entropy_tuning
self.device = torch.device("cuda" if args.cuda else "cpu")
# Q(s,a) 网络
self.critic = QNetwork(num_inputs, action_space.shape[0],
args.hidden_size).to(device=self.device)
self.critic_optim = Adam(self.critic.parameters(), lr=args.lr)
self.critic_target = QNetwork(num_inputs, action_space.shape[0],
args.hidden_size).to(self.device)
hard_update(self.critic_target, self.critic)
if self.policy_type == "Gaussian":
# Target Entropy = −dim(A) (e.g. , -6 for HalfCheetah-v2) as given in the paper
if self.automatic_entropy_tuning is True:
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=args.lr)
self.policy = GaussianPolicy(num_inputs, action_space.shape[0],
args.hidden_size,
action_space).to(self.device)
self.policy_optim = Adam(self.policy.parameters(), lr=args.lr)
else:
self.alpha = 0
self.automatic_entropy_tuning = False
self.policy = DeterministicPolicy(num_inputs,
action_space.shape[0],
args.hidden_size,
action_space).to(self.device)
self.policy_optim = Adam(self.policy.parameters(), lr=args.lr)
def select_action(self, state, evaluate=False):
state = torch.FloatTensor(state).to(self.device).unsqueeze(0) #增加一个维度
if evaluate is False:
action, _, _ = self.policy.sample(state)
else:
_, _, action = self.policy.sample(state)
return action.detach().cpu().numpy()[0]
def update_parameters(self, memory, batch_size, updates):
# Sample a batch from memory
state_batch, action_batch, reward_batch, next_state_batch, mask_batch = memory.sample(
batch_size=batch_size)
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)
with torch.no_grad():
next_state_action, next_state_log_pi, _ = self.policy.sample(
next_state_batch)
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)
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]
pi, log_pi, _ = self.policy.sample(state_batch)
qf1_pi, qf2_pi = self.critic(state_batch, pi)
min_qf_pi = torch.min(qf1_pi, qf2_pi)
policy_loss = ((self.alpha * log_pi) - min_qf_pi).mean(
) # Jπ = 𝔼st∼D,εt∼N[α * logπ(f(εt;st)|st) − Q(st,f(εt;st))]
self.critic_optim.zero_grad()
qf1_loss.backward()
self.critic_optim.step()
self.critic_optim.zero_grad()
qf2_loss.backward()
self.critic_optim.step()
self.policy_optim.zero_grad()
policy_loss.backward()
self.policy_optim.step()
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()
alpha_tlogs = self.alpha.clone() # For TensorboardX logs
else:
alpha_loss = torch.tensor(0.).to(self.device)
alpha_tlogs = torch.tensor(self.alpha) # For TensorboardX logs
if updates % self.target_update_interval == 0:
soft_update(self.critic_target, self.critic, self.tau)
return qf1_loss.item(), qf2_loss.item(), policy_loss.item(
), alpha_loss.item(), alpha_tlogs.item()
# Save model parameters
def save_model(self,
env_name,
suffix="",
actor_path=None,
critic_path=None):
if not os.path.exists('models/'):
os.makedirs('models/')
if actor_path is None:
actor_path = "models/sac_actor_{}_{}".format(env_name, suffix)
if critic_path is None:
critic_path = "models/sac_critic_{}_{}".format(env_name, suffix)
print('Saving models to {} and {}'.format(actor_path, critic_path))
torch.save(self.policy.state_dict(), actor_path)
torch.save(self.critic.state_dict(), critic_path)
# Load model parameters
def load_model(self, actor_path, critic_path):
print('Loading models from {} and {}'.format(actor_path, critic_path))
if actor_path is not None:
self.policy.load_state_dict(torch.load(actor_path))
if critic_path is not None:
self.critic.load_state_dict(torch.load(critic_path))
class SAC(object):
def __init__(self, num_inputs, action_space, args):
self.num_inputs = num_inputs
self.action_space = action_space.shape[0]
self.gamma = args.gamma
self.tau = args.tau
self.alpha = args.alpha
self.policy_type = args.policy
self.target_update_interval = args.target_update_interval
self.automatic_entropy_tuning = args.automatic_entropy_tuning
self.critic = QNetwork(self.num_inputs, self.action_space,
args.hidden_size)
self.critic_optim = Adam(self.critic.parameters(), lr=args.lr)
# Target Entropy = −dim(A) (e.g. , -6 for HalfCheetah-v2) as given in the paper
if self.automatic_entropy_tuning == True:
self.target_entropy = -torch.prod(torch.Tensor(
action_space.shape)).item()
self.log_alpha = torch.zeros(1, requires_grad=True)
self.alpha_optim = Adam([self.log_alpha], lr=args.lr)
else:
pass
if self.policy_type == "Gaussian":
self.policy = GaussianPolicy(self.num_inputs, self.action_space,
args.hidden_size)
self.policy_optim = Adam(self.policy.parameters(), lr=args.lr)
self.value = ValueNetwork(self.num_inputs, args.hidden_size)
self.value_target = ValueNetwork(self.num_inputs, args.hidden_size)
self.value_optim = Adam(self.value.parameters(), lr=args.lr)
hard_update(self.value_target, self.value)
else:
self.policy = DeterministicPolicy(self.num_inputs,
self.action_space,
args.hidden_size)
self.policy_optim = Adam(self.policy.parameters(), lr=args.lr)
self.critic_target = QNetwork(self.num_inputs, self.action_space,
args.hidden_size)
hard_update(self.critic_target, self.critic)
def select_action(self, state, eval=False):
state = torch.FloatTensor(state).unsqueeze(0)
if eval == False:
self.policy.train()
action, _, _, _, _ = self.policy.evaluate(state)
else:
self.policy.eval()
_, _, _, action, _ = self.policy.evaluate(state)
#action = torch.tanh(action)
action = action.detach().cpu().numpy()
return action[0]
def update_parameters(self, state_batch, action_batch, reward_batch,
next_state_batch, mask_batch, updates):
state_batch = torch.FloatTensor(state_batch)
next_state_batch = torch.FloatTensor(next_state_batch)
action_batch = torch.FloatTensor(action_batch)
reward_batch = torch.FloatTensor(reward_batch)
mask_batch = torch.FloatTensor(np.float32(mask_batch))
reward_batch = reward_batch.unsqueeze(
1) # reward_batch = [batch_size, 1]
mask_batch = mask_batch.unsqueeze(1) # mask_batch = [batch_size, 1]
"""
Use two Q-functions to mitigate positive bias in the policy improvement step that is known
to degrade performance of value based methods. Two Q-functions also significantly speed
up training, especially on harder task.
"""
expected_q1_value, expected_q2_value = self.critic(
state_batch, action_batch)
new_action, log_prob, _, mean, log_std = self.policy.evaluate(
state_batch)
if self.automatic_entropy_tuning:
"""
Alpha Loss
"""
alpha_loss = -(self.log_alpha *
(log_prob + self.target_entropy).detach()).mean()
self.alpha_optim.zero_grad()
alpha_loss.backward()
self.alpha_optim.step()
self.alpha = self.log_alpha.exp()
alpha_logs = self.alpha.clone() # For TensorboardX logs
else:
alpha_loss = torch.tensor(0.)
alpha_logs = self.alpha # For TensorboardX logs
if self.policy_type == "Gaussian":
"""
Including a separate function approximator for the soft value can stabilize training.
"""
expected_value = self.value(state_batch)
target_value = self.value_target(next_state_batch)
next_q_value = reward_batch + mask_batch * self.gamma * target_value
else:
"""
There is no need in principle to include a separate function approximator for the state value.
We use a target critic network for deterministic policy and eradicate the value value network completely.
"""
next_state_action, _, _, _, _, = self.policy.evaluate(
next_state_batch)
target_critic_1, target_critic_2 = self.critic_target(
next_state_batch, next_state_action)
target_critic = torch.min(target_critic_1, target_critic_2)
next_q_value = reward_batch + mask_batch * self.gamma * target_critic
"""
Soft Q-function parameters can be trained to minimize the soft Bellman residual
JQ = 𝔼(st,at)~D[0.5(Q1(st,at) - r(st,at) - γ(𝔼st+1~p[V(st+1)]))^2]
∇JQ = ∇Q(st,at)(Q(st,at) - r(st,at) - γV(target)(st+1))
"""
q1_value_loss = F.mse_loss(expected_q1_value, next_q_value.detach())
q2_value_loss = F.mse_loss(expected_q2_value, next_q_value.detach())
q1_new, q2_new = self.critic(state_batch, new_action)
expected_new_q_value = torch.min(q1_new, q2_new)
if self.policy_type == "Gaussian":
"""
Including a separate function approximator for the soft value can stabilize training and is convenient to
train simultaneously with the other networks
Update the V towards the min of two Q-functions in order to reduce overestimation bias from function approximation error.
JV = 𝔼st~D[0.5(V(st) - (𝔼at~π[Qmin(st,at) - α * log π(at|st)]))^2]
∇JV = ∇V(st)(V(st) - Q(st,at) + (α * logπ(at|st)))
"""
next_value = expected_new_q_value - (self.alpha * log_prob)
value_loss = F.mse_loss(expected_value, next_value.detach())
else:
pass
"""
Reparameterization trick is used to get a low variance estimator
f(εt;st) = action sampled from the policy
εt is an input noise vector, sampled from some fixed distribution
Jπ = 𝔼st∼D,εt∼N[α * logπ(f(εt;st)|st) − Q(st,f(εt;st))]
∇Jπ = ∇log π + ([∇at (α * logπ(at|st)) − ∇at Q(st,at)])∇f(εt;st)
"""
policy_loss = ((self.alpha * log_prob) - expected_new_q_value).mean()
# Regularization Loss
mean_loss = 0.001 * mean.pow(2).mean()
std_loss = 0.001 * log_std.pow(2).mean()
policy_loss += mean_loss + std_loss
self.critic_optim.zero_grad()
q1_value_loss.backward()
self.critic_optim.step()
self.critic_optim.zero_grad()
q2_value_loss.backward()
self.critic_optim.step()
if self.policy_type == "Gaussian":
self.value_optim.zero_grad()
value_loss.backward()
self.value_optim.step()
else:
value_loss = torch.tensor(0.)
self.policy_optim.zero_grad()
policy_loss.backward()
self.policy_optim.step()
"""
We update the target weights to match the current value function weights periodically
Update target parameter after every n(args.target_update_interval) updates
"""
if updates % self.target_update_interval == 0 and self.policy_type == "Deterministic":
soft_update(self.critic_target, self.critic, self.tau)
elif updates % self.target_update_interval == 0 and self.policy_type == "Gaussian":
soft_update(self.value_target, self.value, self.tau)
return value_loss.item(), q1_value_loss.item(), q2_value_loss.item(
), policy_loss.item(), alpha_loss.item(), alpha_logs
# Save model parameters
def save_model(self,
env_name,
suffix="",
actor_path=None,
critic_path=None,
value_path=None):
if not os.path.exists('models/'):
os.makedirs('models/')
if actor_path is None:
actor_path = "models/sac_actor_{}_{}".format(env_name, suffix)
if critic_path is None:
critic_path = "models/sac_critic_{}_{}".format(env_name, suffix)
if value_path is None:
value_path = "models/sac_value_{}_{}".format(env_name, suffix)
print('Saving models to {}, {} and {}'.format(actor_path, critic_path,
value_path))
torch.save(self.value.state_dict(), value_path)
torch.save(self.policy.state_dict(), actor_path)
torch.save(self.critic.state_dict(), critic_path)
# Load model parameters
def load_model(self, actor_path, critic_path, value_path):
print('Loading models from {}, {} and {}'.format(
actor_path, critic_path, value_path))
if actor_path is not None:
self.policy.load_state_dict(torch.load(actor_path))
if critic_path is not None:
self.critic.load_state_dict(torch.load(critic_path))
if value_path is not None:
self.value.load_state_dict(torch.load(value_path))
class Controller():
def __init__(self):
self.rate = rospy.Rate(100)
rospy.Subscriber('robot_0/pose', Float32MultiArray,
self.loc_callback_0)
rospy.Subscriber('robot_1/pose', Float32MultiArray,
self.loc_callback_1)
self.states = np.zeros(6)
self.pub0 = rospy.Publisher('/robot_0/cmd_vel', Twist, queue_size=10)
self.pub1 = rospy.Publisher('/robot_1/cmd_vel', Twist, queue_size=10)
# self.listener = TransformListener()
# self.__timer_current = rospy.Timer(rospy.Duration(0.01), self.loc)
self.control_cmd0 = Twist()
self.control_cmd1 = Twist()
self.u_range = np.array([0.5, 2.0])
self.action_space = spaces.Box(low=-self.u_range,
high=+self.u_range,
shape=(2, ),
dtype=np.float32)
self.attacker = DeterministicPolicy(6, 2, 256,
self.action_space).to('cuda')
self.attacker.load_state_dict(torch.load("pretrain"))
self.defender = DeterministicPolicy(6, 2, 256,
self.action_space).to('cuda')
self.defender.load_state_dict(torch.load("d_actor"))
if (control_mode == 1):
self.auto()
else:
self.manual()
def check_goal(self, a_state):
if (0.5 <= a_state[0] <= 1.5 and 0.5 <= a_state[1] <= 1.5):
return True
else:
return False
def _get_done(self):
a_state = self.states[:3]
d_state = self.states[3:]
distance = np.sqrt(np.sum(np.square(a_state[:2] - d_state[:2])))
# print(distance)
if (distance < 0.10 or self.check_goal(a_state)):
done_n = True
else:
done_n = False
if (a_state[1] > 1.4 or a_state[1] < -1.4 or a_state[0] < -1.4
or a_state[0] > 1.4):
done_n = True
if (d_state[1] > 1.4 or d_state[1] < -1.4 or d_state[0] < -1.4
or d_state[0] > 1.4):
done_n = True
return done_n
def loc_callback_0(self, msg):
for i in range(3):
self.states[i] = msg.data[i]
def loc_callback_1(self, msg):
for i in range(3):
self.states[i + 3] = msg.data[i]
def cmd_0(self, data):
self.control_cmd0.linear.x = data[0]
self.control_cmd0.angular.z = data[1]
self.pub0.publish(self.control_cmd0)
def cmd_1(self, data):
self.control_cmd1.linear.x = data[0]
self.control_cmd1.angular.z = data[1]
self.pub1.publish(self.control_cmd1)
def auto(self):
record = []
while not rospy.is_shutdown():
input_a = torch.FloatTensor(
self.states.flatten()).to("cuda").unsqueeze(0)
dxu0 = self.attacker(input_a).detach().cpu().numpy()[0]
input_d = torch.FloatTensor(
self.states.flatten()).to("cuda").unsqueeze(0)
dxu1 = self.defender(input_d).detach().cpu().numpy()[0]
# dxu0 = MPC_controller(self.states[:3],self.states[3:])
# dxu1 = Fast_controller_defence(self.states[:3],self.states[3:])
# dxu1 = Fast_Catch(self.states[:3],self.states[3:])
a = np.copy(self.states)
record.append(a)
# self.cmd_0(dxu0)
# self.cmd_1(dxu1)
if (self._get_done() == True):
dxu1 = np.array([0.0, 0.0])
dxu0 = np.array([0.0, 0.0])
print("stop!")
self.cmd_0(dxu0)
self.cmd_1(dxu1)
np.save("record.npy", record)
# print(dxu0)
self.rate.sleep()
def getKey(self):
if os.name == 'nt':
return msvcrt.getch()
tty.setraw(sys.stdin.fileno())
rlist, _, _ = select.select([sys.stdin], [], [], 0.1)
if rlist:
key = sys.stdin.read(1)
else:
key = ''
termios.tcsetattr(sys.stdin, termios.TCSADRAIN, settings)
return key
def manual(self):
record = []
data = np.array([0.0, 0.0])
data_2 = np.array([0.0, 0.0])
while not rospy.is_shutdown():
print(data, data_2)
# print(record)
a = np.copy(self.states)
record.append(a)
key = self.getKey()
if key == 'w':
if (data[0] < 0.5):
data[0] = data[0] + 0.1
else:
data = data
elif key == 'x':
if (data[0] > -0.5):
data[0] = data[0] - 0.1
else:
data = data
elif key == 'a':
if (data[1] < 2):
data[1] += 0.5
else:
data = data
elif key == 'd':
if (data[1] > -2):
data[1] -= 0.5
else:
data = data
elif key == 's':
data = np.array([0.0, 0.0])
elif (key == '\x03'):
break
else:
data = data
# if(self._get_done()==True):
# data=np.array([ 0.0, 0.0])
# print("stop!")
np.save("record.npy", record)
self.cmd_0(data)
self.rate.sleep()
class SAC(object):
def __init__(self,
num_inputs,
action_space,
args,
process_obs=None,
opt_level='O1'):
self.gamma = args.gamma
self.tau = args.tau
self.alpha = args.alpha
self.device = torch.device("cuda" if args.cuda else "cpu")
self.dtype = torch.float
self.policy_type = args.policy
self.target_update_interval = args.target_update_interval
self.automatic_entropy_tuning = args.automatic_entropy_tuning
self.process_obs = process_obs.to(self.device).to(self.dtype)
self.critic = QNetwork(num_inputs, action_space.shape[0],
args.hidden_size).to(device=self.device).to(
self.dtype)
self.critic_optim = Adam(list(self.critic.parameters()) +
list(process_obs.parameters()),
lr=args.lr)
self.critic_target = QNetwork(num_inputs, action_space.shape[0],
args.hidden_size).to(self.device).to(
self.dtype)
hard_update(self.critic_target, self.critic)
if self.policy_type == "Gaussian":
# Target Entropy = −dim(A) (e.g. , -6 for HalfCheetah-v2) as given in the paper
if self.automatic_entropy_tuning is True:
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,
dtype=self.dtype)
self.alpha_optim = Adam([self.log_alpha], lr=args.lr)
self.policy = GaussianPolicy(num_inputs, action_space.shape[0],
args.hidden_size, action_space).to(
self.device).to(self.dtype)
self.policy_optim = Adam(list(self.policy.parameters()) +
list(process_obs.parameters()),
lr=args.lr)
else:
self.alpha = 0
self.automatic_entropy_tuning = False
self.policy = DeterministicPolicy(
num_inputs, action_space.shape[0], args.hidden_size,
action_space).to(self.device).to(self.dtype)
self.policy_optim = Adam(list(self.policy.parameters()) +
list(process_obs.parameters()),
lr=args.lr)
if opt_level is not None:
model, optimizer = amp.initialize([
self.policy, self.process_obs, self.critic, self.critic_target
], [self.policy_optim, self.critic_optim],
opt_level=opt_level)
def select_action(self, obs, evaluate=False):
with torch.no_grad():
obs = torch.FloatTensor(obs).to(self.device).unsqueeze(0).to(
self.dtype)
state = self.process_obs(obs)
if evaluate is False:
action, _, _ = self.policy.sample(state)
else:
_, _, action = self.policy.sample(state)
action = action.detach().cpu().numpy()[0]
return action
def update_parameters(self, memory, batch_size, updates):
# Sample a batch from memory
obs_batch, action_batch, reward_batch, next_obs_batch, mask_batch = memory.sample(
batch_size=batch_size)
obs_batch = torch.FloatTensor(obs_batch).to(self.device).to(self.dtype)
next_obs_batch = torch.FloatTensor(next_obs_batch).to(self.device).to(
self.dtype)
action_batch = torch.FloatTensor(action_batch).to(self.device).to(
self.dtype)
reward_batch = torch.FloatTensor(reward_batch).to(
self.device).unsqueeze(1).to(self.dtype)
mask_batch = torch.FloatTensor(mask_batch).to(
self.device).unsqueeze(1).to(self.dtype)
state_batch = self.process_obs(obs_batch)
with torch.no_grad():
next_state_batch = self.process_obs(next_obs_batch)
next_state_action, next_state_log_pi, _ = self.policy.sample(
next_state_batch)
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)
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
self.critic_optim.zero_grad()
assert torch.isfinite(qf_loss).all()
with amp.scale_loss(qf_loss, self.critic_optim) as qf_loss:
qf_loss.backward()
self.critic_optim.step()
state_batch = self.process_obs(obs_batch)
pi, log_pi, _ = self.policy.sample(state_batch)
qf1_pi, qf2_pi = self.critic(state_batch.detach(), pi)
min_qf_pi = torch.min(qf1_pi, qf2_pi)
policy_loss = ((self.alpha * log_pi) - min_qf_pi).mean(
) # Jπ = 𝔼st∼D,εt∼N[α * logπ(f(εt;st)|st) − Q(st,f(εt;st))]
self.policy_optim.zero_grad()
assert torch.isfinite(policy_loss).all()
with amp.scale_loss(policy_loss, self.policy_optim) as policy_loss:
policy_loss.backward()
self.policy_optim.step()
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()
alpha_tlogs = self.alpha.clone() # For TensorboardX logs
else:
alpha_loss = torch.tensor(0.).to(self.device).to(self.dtype)
alpha_tlogs = torch.tensor(self.alpha) # For TensorboardX logs
if updates % self.target_update_interval == 0:
soft_update(self.critic_target, self.critic, self.tau)
return qf1_loss.item(), qf2_loss.item(), policy_loss.item(
), alpha_loss.item(), alpha_tlogs.item()
# Save model parameters
def save_model(self,
actor_path=None,
critic_path=None,
process_obs_path=None):
logger.debug(
f'saving models to {actor_path} and {critic_path} and {process_obs_path}'
)
torch.save(self.policy.state_dict(), actor_path)
torch.save(self.critic.state_dict(), critic_path)
torch.save(self.process_obs.state_dict(), process_obs_path)
# Load model parameters
def load_model(self,
actor_path=None,
critic_path=None,
process_obs_path=None):
logger.info(
f'Loading models from {actor_path} and {critic_path} and {process_obs_path}'
)
if actor_path is not None:
self.policy.load_state_dict(torch.load(actor_path))
if critic_path is not None:
self.critic.load_state_dict(torch.load(critic_path))
if process_obs_path is not None:
self.process_obs.load_state_dict(torch.load(process_obs_path))
class SAC(object):
def __init__(self, num_inputs, action_space, args):
self.gamma = args.gamma
self.tau = args.tau
self.alpha = args.alpha
self.policy_type = args.policy_type
self.target_update_interval = args.target_update_interval
self.automatic_entropy_tuning = args.automatic_entropy_tuning
self.device = torch.device("cuda" if args.cuda else "cpu")
self.critic = QNetwork(num_inputs, action_space.shape[0],
args.hidden_size).to(device=self.device)
self.critic_optim = Adam(self.critic.parameters(), lr=args.lr)
self.critic_target = QNetwork(num_inputs, action_space.shape[0],
args.hidden_size).to(self.device)
hard_update(self.critic_target, self.critic)
self.obs_upper_bound, self.obs_lower_bound = 0, 0
self.reward_lower_bound, self.reward_upper_bound = 0, 0
if self.policy_type == "Gaussian":
# Target Entropy = −dim(A) (e.g. , -6 for HalfCheetah-v2) as given in the paper
if self.automatic_entropy_tuning is True:
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=args.lr)
self.policy = GaussianPolicy(num_inputs, action_space.shape[0],
args.hidden_size,
action_space).to(self.device)
self.policy_optim = Adam(self.policy.parameters(), lr=args.lr)
else:
self.alpha = 0
self.automatic_entropy_tuning = False
self.policy = DeterministicPolicy(num_inputs,
action_space.shape[0],
args.hidden_size,
action_space).to(self.device)
self.policy_optim = Adam(self.policy.parameters(), lr=args.lr)
def select_action(self, state, evaluate=False):
state = torch.FloatTensor(state).to(self.device).unsqueeze(0)
if evaluate is False:
action, _, _ = self.policy.sample(state)
else:
_, _, action = self.policy.sample(state)
return action.detach().cpu().numpy()[0]
def update_parameters(self, memory, batch_size, updates):
# Sample a batch from memory
state_batch, action_batch, reward_batch, next_state_batch, mask_batch = memory.sample(
batch_size=batch_size)
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)
with torch.no_grad():
next_state_action, next_state_log_pi, _ = self.policy.sample(
next_state_batch)
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)
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
self.critic_optim.zero_grad()
qf_loss.backward()
self.critic_optim.step()
pi, log_pi, _ = self.policy.sample(state_batch)
qf1_pi, qf2_pi = self.critic(state_batch, pi)
min_qf_pi = torch.min(qf1_pi, qf2_pi)
policy_loss = ((self.alpha * log_pi) - min_qf_pi).mean(
) # Jπ = 𝔼st∼D,εt∼N[α * logπ(f(εt;st)|st) − Q(st,f(εt;st))]
self.policy_optim.zero_grad()
policy_loss.backward()
self.policy_optim.step()
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()
alpha_tlogs = self.alpha.clone() # For TensorboardX logs
else:
alpha_loss = torch.tensor(0.).to(self.device)
alpha_tlogs = torch.tensor(self.alpha) # For TensorboardX logs
if updates % self.target_update_interval == 0:
soft_update(self.critic_target, self.critic, self.tau)
return qf1_loss.item(), qf2_loss.item(), policy_loss.item(
), alpha_loss.item(), alpha_tlogs.item()
# Save model parameters
def save(self, filename):
torch.save(self.critic.state_dict(), filename + "_critic")
torch.save(self.critic_optim.state_dict(),
filename + "_critic_optimizer")
torch.save(self.policy.state_dict(), filename + "_actor")
torch.save(self.policy_optim.state_dict(),
filename + "_actor_optimizer")
def load(self, filename):
self.critic.load_state_dict(torch.load(filename + "_critic.pth"))
self.critic_optim.load_state_dict(
torch.load(filename + "_critic_optimizer"))
self.critic_target = copy.deepcopy(self.critic)
self.policy.load_state_dict(torch.load(filename + "_actor.pth"))
self.policy_optim.load_state_dict(
torch.load(filename + "_actor_optimizer"))
