class IOC(Off_Policy):
'''
Learning Options with Interest Functions, https://www.aaai.org/ojs/index.php/AAAI/article/view/5114/4987
Options of Interest: Temporal Abstraction with Interest Functions, http://arxiv.org/abs/2001.00271
'''
def __init__(
self,
envspec,
q_lr=5.0e-3,
intra_option_lr=5.0e-4,
termination_lr=5.0e-4,
interest_lr=5.0e-4,
boltzmann_temperature=1.0,
options_num=4,
ent_coff=0.01,
double_q=False,
use_baseline=True,
terminal_mask=True,
termination_regularizer=0.01,
assign_interval=1000,
network_settings={
'q': [32, 32],
'intra_option': [32, 32],
'termination': [32, 32],
'interest': [32, 32]
},
**kwargs):
super().__init__(envspec=envspec, **kwargs)
self.assign_interval = assign_interval
self.options_num = options_num
self.termination_regularizer = termination_regularizer
self.ent_coff = ent_coff
self.use_baseline = use_baseline
self.terminal_mask = terminal_mask
self.double_q = double_q
self.boltzmann_temperature = boltzmann_temperature
def _create_net(name, representation_net=None):
return ValueNetwork(
name=name,
representation_net=representation_net,
value_net_type=OutputNetworkType.CRITIC_QVALUE_ALL,
value_net_kwargs=dict(output_shape=self.options_num,
network_settings=network_settings['q']))
self.q_net = _create_net('q_net', self._representation_net)
self._representation_target_net = self._create_representation_net(
'_representation_target_net')
self.q_target_net = _create_net('q_target_net',
self._representation_target_net)
self.intra_option_net = ValueNetwork(
name='intra_option_net',
value_net_type=OutputNetworkType.OC_INTRA_OPTION,
value_net_kwargs=dict(
vector_dim=self._representation_net.h_dim,
output_shape=self.a_dim,
options_num=self.options_num,
network_settings=network_settings['intra_option']))
self.termination_net = ValueNetwork(
name='termination_net',
value_net_type=OutputNetworkType.CRITIC_QVALUE_ALL,
value_net_kwargs=dict(
vector_dim=self._representation_net.h_dim,
output_shape=self.options_num,
network_settings=network_settings['termination'],
out_activation='sigmoid'))
self.interest_net = ValueNetwork(
name='interest_net',
value_net_type=OutputNetworkType.CRITIC_QVALUE_ALL,
value_net_kwargs=dict(
vector_dim=self._representation_net.h_dim,
output_shape=self.options_num,
network_settings=network_settings['interest'],
out_activation='sigmoid'))
self.actor_tv = self.intra_option_net.trainable_variables
if self.is_continuous:
self.log_std = tf.Variable(initial_value=-0.5 * np.ones(
(self.options_num, self.a_dim), dtype=np.float32),
trainable=True) # [P, A]
self.actor_tv += [self.log_std]
update_target_net_weights(self.q_target_net.weights,
self.q_net.weights)
self.q_lr, self.intra_option_lr, self.termination_lr, self.interest_lr = map(
self.init_lr, [q_lr, intra_option_lr, termination_lr, interest_lr])
self.q_optimizer = self.init_optimizer(self.q_lr, clipvalue=5.)
self.intra_option_optimizer = self.init_optimizer(self.intra_option_lr,
clipvalue=5.)
self.termination_optimizer = self.init_optimizer(self.termination_lr,
clipvalue=5.)
self.interest_optimizer = self.init_optimizer(self.interest_lr,
clipvalue=5.)
self._worker_params_dict.update(self.q_net._policy_models)
self._worker_params_dict.update(self.intra_option_net._policy_models)
self._worker_params_dict.update(self.interest_net._policy_models)
self._all_params_dict.update(self.q_net._all_models)
self._all_params_dict.update(self.intra_option_net._all_models)
self._all_params_dict.update(self.interest_net._all_models)
self._all_params_dict.update(self.termination_net._all_models)
self._all_params_dict.update(
q_optimizer=self.q_optimizer,
intra_option_optimizer=self.intra_option_optimizer,
termination_optimizer=self.termination_optimizer,
interest_optimizer=self.interest_optimizer)
self._model_post_process()
def _generate_random_options(self):
return tf.constant(np.random.randint(0, self.options_num,
self.n_agents),
dtype=tf.int32)
def choose_action(self, s, visual_s, evaluation=False):
if not hasattr(self, 'options'):
self.options = self._generate_random_options()
self.last_options = self.options
a, self.options, self.cell_state = self._get_action(
s, visual_s, self.cell_state, self.options)
a = a.numpy()
return a
@tf.function
def _get_action(self, s, visual_s, cell_state, options):
with tf.device(self.device):
feat, cell_state = self._representation_net(s,
visual_s,
cell_state=cell_state)
q = self.q_net.value_net(feat) # [B, P]
pi = self.intra_option_net.value_net(feat) # [B, P, A]
options_onehot = tf.one_hot(options,
self.options_num,
dtype=tf.float32) # [B, P]
options_onehot_expanded = tf.expand_dims(options_onehot,
axis=-1) # [B, P, 1]
pi = tf.reduce_sum(pi * options_onehot_expanded, axis=1) # [B, A]
if self.is_continuous:
log_std = tf.gather(self.log_std, options)
mu = tf.math.tanh(pi)
a, _ = gaussian_clip_rsample(mu, log_std)
else:
pi = pi / self.boltzmann_temperature
dist = tfp.distributions.Categorical(
logits=tf.nn.log_softmax(pi)) # [B, ]
a = dist.sample()
interests = self.interest_net.value_net(feat) # [B, P]
op_logits = interests * q # [B, P] or tf.nn.softmax(q)
new_options = tfp.distributions.Categorical(
logits=tf.nn.log_softmax(op_logits)).sample()
return a, new_options, cell_state
def _target_params_update(self):
if self.global_step % self.assign_interval == 0:
update_target_net_weights(self.q_target_net.weights,
self.q_net.weights)
def learn(self, **kwargs):
self.train_step = kwargs.get('train_step')
for i in range(self.train_times_per_step):
self._learn(
function_dict={
'sample_data_list': [
's', 'visual_s', 'a', 'r', 's_', 'visual_s_', 'done',
'last_options', 'options'
],
'train_data_list': [
's', 'visual_s', 'a', 'r', 's_', 'visual_s_', 'done',
'last_options', 'options'
],
'summary_dict':
dict([['LEARNING_RATE/q_lr',
self.q_lr(self.train_step)],
[
'LEARNING_RATE/intra_option_lr',
self.intra_option_lr(self.train_step)
],
[
'LEARNING_RATE/termination_lr',
self.termination_lr(self.train_step)
], ['Statistics/option', self.options[0]]])
})
@tf.function(experimental_relax_shapes=True)
def _train(self, memories, isw, cell_state):
s, visual_s, a, r, s_, visual_s_, done, last_options, options = memories
last_options = tf.cast(last_options, tf.int32)
options = tf.cast(options, tf.int32)
with tf.device(self.device):
with tf.GradientTape(persistent=True) as tape:
feat, _ = self._representation_net(s,
visual_s,
cell_state=cell_state)
feat_, _ = self._representation_target_net(
s_, visual_s_, cell_state=cell_state)
q = self.q_net.value_net(feat) # [B, P]
pi = self.intra_option_net.value_net(feat) # [B, P, A]
beta = self.termination_net.value_net(feat) # [B, P]
q_next = self.q_target_net.value_net(
feat_) # [B, P], [B, P, A], [B, P]
beta_next = self.termination_net.value_net(feat_) # [B, P]
interests = self.interest_net.value_net(feat) # [B, P]
options_onehot = tf.one_hot(options,
self.options_num,
dtype=tf.float32) # [B,] => [B, P]
q_s = qu_eval = tf.reduce_sum(q * options_onehot,
axis=-1,
keepdims=True) # [B, 1]
beta_s_ = tf.reduce_sum(beta_next * options_onehot,
axis=-1,
keepdims=True) # [B, 1]
q_s_ = tf.reduce_sum(q_next * options_onehot,
axis=-1,
keepdims=True) # [B, 1]
if self.double_q:
q_ = self.q_net.value_net(
feat) # [B, P], [B, P, A], [B, P]
max_a_idx = tf.one_hot(
tf.argmax(q_, axis=-1),
self.options_num,
dtype=tf.float32) # [B, P] => [B, ] => [B, P]
q_s_max = tf.reduce_sum(q_next * max_a_idx,
axis=-1,
keepdims=True) # [B, 1]
else:
q_s_max = tf.reduce_max(q_next, axis=-1,
keepdims=True) # [B, 1]
u_target = (1 - beta_s_) * q_s_ + beta_s_ * q_s_max # [B, 1]
qu_target = tf.stop_gradient(r + self.gamma *
(1 - done) * u_target)
td_error = qu_target - qu_eval # gradient : q
q_loss = tf.reduce_mean(tf.square(td_error) *
isw) # [B, 1] => 1
if self.use_baseline:
adv = tf.stop_gradient(qu_target - qu_eval)
else:
adv = tf.stop_gradient(qu_target)
options_onehot_expanded = tf.expand_dims(
options_onehot, axis=-1) # [B, P] => [B, P, 1]
pi = tf.reduce_sum(pi * options_onehot_expanded,
axis=1) # [B, P, A] => [B, A]
if self.is_continuous:
log_std = tf.gather(self.log_std, options)
mu = tf.math.tanh(pi)
log_p = gaussian_likelihood_sum(a, mu, log_std)
entropy = gaussian_entropy(log_std)
else:
pi = pi / self.boltzmann_temperature
log_pi = tf.nn.log_softmax(pi, axis=-1) # [B, A]
entropy = -tf.reduce_sum(tf.exp(log_pi) * log_pi,
axis=1,
keepdims=True) # [B, 1]
log_p = tf.reduce_sum(a * log_pi, axis=-1,
keepdims=True) # [B, 1]
pi_loss = tf.reduce_mean(
-(log_p * adv + self.ent_coff * entropy)
) # [B, 1] * [B, 1] => [B, 1] => 1
last_options_onehot = tf.one_hot(
last_options, self.options_num,
dtype=tf.float32) # [B,] => [B, P]
beta_s = tf.reduce_sum(beta * last_options_onehot,
axis=-1,
keepdims=True) # [B, 1]
pi_op = tf.nn.softmax(
interests *
tf.stop_gradient(q)) # [B, P] or tf.nn.softmax(q)
interest_loss = -tf.reduce_mean(beta_s * tf.reduce_sum(
pi_op * options_onehot, axis=-1, keepdims=True) *
q_s) # [B, 1] => 1
v_s = tf.reduce_sum(q * pi_op, axis=-1,
keepdims=True) # [B, P] * [B, P] => [B, 1]
beta_loss = beta_s * tf.stop_gradient(q_s - v_s) # [B, 1]
if self.terminal_mask:
beta_loss *= (1 - done)
beta_loss = tf.reduce_mean(beta_loss) # [B, 1] => 1
q_grads = tape.gradient(q_loss, self.q_net.trainable_variables)
intra_option_grads = tape.gradient(pi_loss, self.actor_tv)
termination_grads = tape.gradient(
beta_loss, self.termination_net.trainable_variables)
interest_grads = tape.gradient(
interest_loss, self.interest_net.trainable_variables)
self.q_optimizer.apply_gradients(
zip(q_grads, self.q_net.trainable_variables))
self.intra_option_optimizer.apply_gradients(
zip(intra_option_grads, self.actor_tv))
self.termination_optimizer.apply_gradients(
zip(termination_grads,
self.termination_net.trainable_variables))
self.interest_optimizer.apply_gradients(
zip(interest_grads, self.interest_net.trainable_variables))
self.global_step.assign_add(1)
return td_error, dict(
[['LOSS/q_loss', tf.reduce_mean(q_loss)],
['LOSS/pi_loss', tf.reduce_mean(pi_loss)],
['LOSS/beta_loss',
tf.reduce_mean(beta_loss)],
['LOSS/interest_loss',
tf.reduce_mean(interest_loss)],
['Statistics/q_option_max',
tf.reduce_max(q_s)],
['Statistics/q_option_min',
tf.reduce_min(q_s)],
['Statistics/q_option_mean',
tf.reduce_mean(q_s)]])
def store_data(self, s, visual_s, a, r, s_, visual_s_, done):
"""
for off-policy training, use this function to store <s, a, r, s_, done> into ReplayBuffer.
"""
assert isinstance(a,
np.ndarray), "store need action type is np.ndarray"
assert isinstance(r,
np.ndarray), "store need reward type is np.ndarray"
assert isinstance(done,
np.ndarray), "store need done type is np.ndarray"
self._running_average(s)
self.data.add(
s,
visual_s,
a,
r[:, np.newaxis], # 升维
s_,
visual_s_,
done[:, np.newaxis], # 升维
self.last_options,
self.options)
def no_op_store(self, s, visual_s, a, r, s_, visual_s_, done):
pass
class CURL(Off_Policy):
"""
CURL: Contrastive Unsupervised Representations for Reinforcement Learning, http://arxiv.org/abs/2004.04136
"""
def __init__(
self,
envspec,
alpha=0.2,
annealing=True,
last_alpha=0.01,
ployak=0.995,
discrete_tau=1.0,
network_settings={
'actor_continuous': {
'share': [128, 128],
'mu': [64],
'log_std': [64],
'soft_clip': False,
'log_std_bound': [-20, 2]
},
'actor_discrete': [64, 32],
'q': [128, 128],
'encoder': 128
},
auto_adaption=True,
actor_lr=5.0e-4,
critic_lr=1.0e-3,
alpha_lr=5.0e-4,
curl_lr=5.0e-4,
img_size=64,
**kwargs):
super().__init__(envspec=envspec, **kwargs)
self.concat_vector_dim = self.obs_spec.total_vector_dim
self.ployak = ployak
self.discrete_tau = discrete_tau
self.auto_adaption = auto_adaption
self.annealing = annealing
self.img_size = img_size
self.img_dim = [img_size, img_size, self.obs_spec.visual_dims[0][-1]]
self.vis_feat_size = network_settings['encoder']
if self.auto_adaption:
self.log_alpha = tf.Variable(initial_value=0.0,
name='log_alpha',
dtype=tf.float32,
trainable=True)
else:
self.log_alpha = tf.Variable(initial_value=tf.math.log(alpha),
name='log_alpha',
dtype=tf.float32,
trainable=False)
if self.annealing:
self.alpha_annealing = LinearAnnealing(alpha, last_alpha,
1.0e6)
def _create_net(name):
return DoubleValueNetwork(
name=name,
value_net_type=OutputNetworkType.CRITIC_QVALUE_ONE,
value_net_kwargs=dict(vector_dim=self.concat_vector_dim +
self.vis_feat_size,
action_dim=self.a_dim,
network_settings=network_settings['q']))
self.critic_net = _create_net('critic_net')
self.critic_target_net = _create_net('critic_target_net')
if self.is_continuous:
self.actor_net = ValueNetwork(
name='actor_net',
value_net_type=OutputNetworkType.ACTOR_CTS,
value_net_kwargs=dict(
vector_dim=self.concat_vector_dim + self.vis_feat_size,
output_shape=self.a_dim,
network_settings=network_settings['actor_continuous']))
else:
self.actor_net = ValueNetwork(
name='actor_net',
value_net_type=OutputNetworkType.ACTOR_DCT,
value_net_kwargs=dict(
vector_dim=self.concat_vector_dim + self.vis_feat_size,
output_shape=self.a_dim,
network_settings=network_settings['actor_discrete']))
self.gumbel_dist = tfp.distributions.Gumbel(0, 1)
# entropy = -log(1/|A|) = log |A|
self.target_entropy = 0.98 * (-self.a_dim if self.is_continuous else
np.log(self.a_dim))
self.encoder = VisualEncoder(self.img_dim, self.vis_feat_size)
self.encoder_target = VisualEncoder(self.img_dim, self.vis_feat_size)
self.curl_w = tf.Variable(
initial_value=tf.random.normal(shape=(self.vis_feat_size,
self.vis_feat_size)),
name='curl_w',
dtype=tf.float32,
trainable=True)
self.critic_tv = self.critic_net.trainable_variables + self.encoder.trainable_variables
update_target_net_weights(
self.critic_target_net.weights +
self.encoder_target.trainable_variables,
self.critic_net.weights + self.encoder.trainable_variables)
self.actor_lr, self.critic_lr, self.alpha_lr, self.curl_lr = map(
self.init_lr, [actor_lr, critic_lr, alpha_lr, curl_lr])
self.optimizer_actor, self.optimizer_critic, self.optimizer_alpha, self.optimizer_curl = map(
self.init_optimizer,
[self.actor_lr, self.critic_lr, self.alpha_lr, self.curl_lr])
self._worker_params_dict.update(self.actor_net._policy_models)
self._worker_params_dict.update(encoder=self.encoder)
self._all_params_dict.update(self.actor_net._all_models)
self._all_params_dict.update(self.critic_net._all_models)
self._all_params_dict.update(curl_w=self.curl_w,
encoder=self.encoder,
optimizer_actor=self.optimizer_actor,
optimizer_critic=self.optimizer_critic,
optimizer_alpha=self.optimizer_alpha,
optimizer_curl=self.optimizer_curl)
self._model_post_process()
def choose_action(self, obs, evaluation=False):
visual = center_crop_image(obs.first_visual()[:, 0], self.img_size)
mu, pi = self._get_action(visual)
a = mu.numpy() if evaluation else pi.numpy()
return a
@tf.function
def _get_action(self, visual):
with tf.device(self.device):
feat = tf.concat([self.encoder(visual),
obs.flatten_vector()],
axis=-1)
if self.is_continuous:
mu, log_std = self.actor_net.value_net(feat)
pi, _ = squash_rsample(mu, log_std)
mu = tf.tanh(mu) # squash mu
else:
logits = self.actor_net.value_net(feat)
mu = tf.argmax(logits, axis=1)
cate_dist = tfp.distributions.Categorical(
logits=tf.nn.log_softmax(logits))
pi = cate_dist.sample()
return mu, pi
def _process_before_train(self, data: BatchExperiences):
visual = np.transpose(data.obs.first_visual()[:, 0].numpy(),
(0, 3, 1, 2))
visual_ = np.transpose(data.obs_.first_visual()[:, 0].numpy(),
(0, 3, 1, 2))
pos = np.transpose(random_crop(visual, self.img_size), (0, 2, 3, 1))
visual = np.transpose(random_crop(visual, self.img_size), (0, 2, 3, 1))
visual_ = np.transpose(random_crop(visual_, self.img_size),
(0, 2, 3, 1))
return self.data_convert([visual, visual_, pos])
def _target_params_update(self):
update_target_net_weights(
self.critic_target_net.weights +
self.encoder_target.trainable_variables,
self.critic_net.weights + self.encoder.trainable_variables,
self.ployak)
def learn(self, **kwargs):
self.train_step = kwargs.get('train_step')
for i in range(self.train_times_per_step):
self._learn(
function_dict={
'summary_dict':
dict([[
'LEARNING_RATE/actor_lr',
self.actor_lr(self.train_step)
],
[
'LEARNING_RATE/critic_lr',
self.critic_lr(self.train_step)
],
[
'LEARNING_RATE/alpha_lr',
self.alpha_lr(self.train_step)
]])
})
@property
def alpha(self):
return tf.exp(self.log_alpha)
def _train(self, BATCH: BatchExperiences, isw, cell_state):
visual, visual_, pos = self._process_before_train(BATCH)
td_error, summaries = self.train(BATCH, isw, cell_state, visual,
visual_, pos)
if self.annealing and not self.auto_adaption:
self.log_alpha.assign(
tf.math.log(
tf.cast(self.alpha_annealing(self.global_step.numpy()),
tf.float32)))
return td_error, summaries
@tf.function
def train(self, BATCH, isw, cell_state, visual, visual_, pos):
with tf.device(self.device):
with tf.GradientTape(persistent=True) as tape:
vis_feat = self.encoder(visual)
vis_feat_ = self.encoder(visual_)
target_vis_feat_ = self.encoder_target(visual_)
feat = tf.concat(
[vis_feat, BATCH.obs.flatten_vector()], axis=-1)
feat_ = tf.concat(
[vis_feat_, BATCH.obs_.flatten_vector()], axis=-1)
target_feat_ = tf.concat(
[target_vis_feat_,
BATCH.obs_.flatten_vector()], axis=-1)
if self.is_continuous:
target_mu, target_log_std = self.actor_net.value_net(feat_)
target_pi, target_log_pi = squash_rsample(
target_mu, target_log_std)
else:
target_logits = self.actor_net.value_net(feat_)
target_cate_dist = tfp.distributions.Categorical(
logits=tf.nn.log_softmax(target_logits))
target_pi = target_cate_dist.sample()
target_log_pi = target_cate_dist.log_prob(target_pi)
target_pi = tf.one_hot(target_pi,
self.a_dim,
dtype=tf.float32)
q1, q2 = self.critic_net.value_net(feat, BATCH.action)
q1_target, q2_target = self.critic_target_net.value_net(
feat_, target_pi)
q_target = tf.minimum(q1_target, q2_target)
dc_r = tf.stop_gradient(
BATCH.reward + self.gamma * (1 - BATCH.done) *
(q_target - self.alpha * target_log_pi))
td_error1 = q1 - dc_r
td_error2 = q2 - dc_r
q1_loss = tf.reduce_mean(tf.square(td_error1) * isw)
q2_loss = tf.reduce_mean(tf.square(td_error2) * isw)
critic_loss = 0.5 * q1_loss + 0.5 * q2_loss
z_a = vis_feat # [B, N]
z_out = self.encoder_target(pos)
logits = tf.matmul(
z_a, tf.matmul(self.curl_w, tf.transpose(z_out, [1, 0])))
logits -= tf.reduce_max(logits, axis=-1, keepdims=True)
curl_loss = tf.reduce_mean(
tf.keras.losses.sparse_categorical_crossentropy(
tf.range(self.batch_size), logits))
critic_grads = tape.gradient(critic_loss, self.critic_tv)
self.optimizer_critic.apply_gradients(
zip(critic_grads, self.critic_tv))
curl_grads = tape.gradient(curl_loss, [self.curl_w] +
self.encoder.trainable_variables)
self.optimizer_curl.apply_gradients(
zip(curl_grads,
[self.curl_w] + self.encoder.trainable_variables))
with tf.GradientTape() as tape:
if self.is_continuous:
mu, log_std = self.actor_net.value_net(feat)
pi, log_pi = squash_rsample(mu, log_std)
entropy = gaussian_entropy(log_std)
else:
logits = self.actor_net.value_net(feat)
logp_all = tf.nn.log_softmax(logits)
gumbel_noise = tf.cast(self.gumbel_dist.sample(
BATCH.action.shape),
dtype=tf.float32)
_pi = tf.nn.softmax(
(logp_all + gumbel_noise) / self.discrete_tau)
_pi_true_one_hot = tf.one_hot(tf.argmax(_pi, axis=-1),
self.a_dim)
_pi_diff = tf.stop_gradient(_pi_true_one_hot - _pi)
pi = _pi_diff + _pi
log_pi = tf.reduce_sum(tf.multiply(logp_all, pi),
axis=1,
keepdims=True)
entropy = -tf.reduce_mean(
tf.reduce_sum(tf.exp(logp_all) * logp_all,
axis=1,
keepdims=True))
q_s_pi = self.critic_net.get_min(feat, pi)
actor_loss = -tf.reduce_mean(q_s_pi - self.alpha * log_pi)
actor_grads = tape.gradient(actor_loss,
self.actor_net.trainable_variables)
self.optimizer_actor.apply_gradients(
zip(actor_grads, self.actor_net.trainable_variables))
if self.auto_adaption:
with tf.GradientTape() as tape:
if self.is_continuous:
mu, log_std = self.actor_net.value_net(feat)
norm_dist = tfp.distributions.Normal(
loc=mu, scale=tf.exp(log_std))
log_pi = tf.reduce_sum(norm_dist.log_prob(
norm_dist.sample()),
axis=-1,
keep_dims=True) # [B, 1]
else:
logits = self.actor_net.value_net(feat)
norm_dist = tfp.distributions.Categorical(
logits=tf.nn.log_softmax(logits))
log_pi = norm_dist.log_prob(cate_dist.sample())
alpha_loss = -tf.reduce_mean(
self.alpha *
tf.stop_gradient(log_pi + self.target_entropy))
alpha_grad = tape.gradient(alpha_loss, self.log_alpha)
self.optimizer_alpha.apply_gradients([(alpha_grad,
self.log_alpha)])
self.global_step.assign_add(1)
summaries = dict(
[['LOSS/actor_loss', actor_loss], ['LOSS/q1_loss', q1_loss],
['LOSS/q2_loss', q2_loss], ['LOSS/critic_loss', critic_loss],
['LOSS/curl_loss', curl_loss],
['Statistics/log_alpha', self.log_alpha],
['Statistics/alpha', self.alpha],
['Statistics/entropy', entropy],
['Statistics/q_min',
tf.reduce_min(tf.minimum(q1, q2))],
['Statistics/q_mean',
tf.reduce_mean(tf.minimum(q1, q2))],
['Statistics/q_max',
tf.reduce_max(tf.maximum(q1, q2))]])
if self.auto_adaption:
summaries.update({'LOSS/alpha_loss': alpha_loss})
return (td_error1 + td_error2) / 2., summaries
class OC(Off_Policy):
'''
The Option-Critic Architecture. http://arxiv.org/abs/1609.05140
'''
def __init__(self,
envspec,
q_lr=5.0e-3,
intra_option_lr=5.0e-4,
termination_lr=5.0e-4,
use_eps_greedy=False,
eps_init=1,
eps_mid=0.2,
eps_final=0.01,
init2mid_annealing_step=1000,
boltzmann_temperature=1.0,
options_num=4,
ent_coff=0.01,
double_q=False,
use_baseline=True,
terminal_mask=True,
termination_regularizer=0.01,
assign_interval=1000,
network_settings={
'q': [32, 32],
'intra_option': [32, 32],
'termination': [32, 32]
},
**kwargs):
super().__init__(envspec=envspec, **kwargs)
self.expl_expt_mng = ExplorationExploitationClass(
eps_init=eps_init,
eps_mid=eps_mid,
eps_final=eps_final,
init2mid_annealing_step=init2mid_annealing_step,
max_step=self.max_train_step)
self.assign_interval = assign_interval
self.options_num = options_num
self.termination_regularizer = termination_regularizer
self.ent_coff = ent_coff
self.use_baseline = use_baseline
self.terminal_mask = terminal_mask
self.double_q = double_q
self.boltzmann_temperature = boltzmann_temperature
self.use_eps_greedy = use_eps_greedy
def _create_net(name, representation_net=None):
return ValueNetwork(
name=name,
representation_net=representation_net,
value_net_type=OutputNetworkType.CRITIC_QVALUE_ALL,
value_net_kwargs=dict(output_shape=self.options_num,
network_settings=network_settings['q']))
self.q_net = _create_net('q_net', self._representation_net)
self._representation_target_net = self._create_representation_net(
'_representation_target_net')
self.q_target_net = _create_net('q_target_net',
self._representation_target_net)
self.intra_option_net = ValueNetwork(
name='intra_option_net',
value_net_type=OutputNetworkType.OC_INTRA_OPTION,
value_net_kwargs=dict(
vector_dim=self._representation_net.h_dim,
output_shape=self.a_dim,
options_num=self.options_num,
network_settings=network_settings['intra_option']))
self.termination_net = ValueNetwork(
name='termination_net',
value_net_type=OutputNetworkType.CRITIC_QVALUE_ALL,
value_net_kwargs=dict(
vector_dim=self._representation_net.h_dim,
output_shape=self.options_num,
network_settings=network_settings['termination'],
out_activation='sigmoid'))
self.actor_tv = self.intra_option_net.trainable_variables
if self.is_continuous:
self.log_std = tf.Variable(initial_value=-0.5 * np.ones(
(self.options_num, self.a_dim), dtype=np.float32),
trainable=True) # [P, A]
self.actor_tv += [self.log_std]
update_target_net_weights(self.q_target_net.weights,
self.q_net.weights)
self.q_lr, self.intra_option_lr, self.termination_lr = map(
self.init_lr, [q_lr, intra_option_lr, termination_lr])
self.q_optimizer = self.init_optimizer(self.q_lr, clipvalue=5.)
self.intra_option_optimizer = self.init_optimizer(self.intra_option_lr,
clipvalue=5.)
self.termination_optimizer = self.init_optimizer(self.termination_lr,
clipvalue=5.)
self._worker_params_dict.update(self.q_net._policy_models)
self._worker_params_dict.update(self.intra_option_net._policy_models)
self._worker_params_dict.update(self.termination_net._policy_models)
self._all_params_dict.update(self.q_net._all_models)
self._all_params_dict.update(self.intra_option_net._all_models)
self._all_params_dict.update(self.termination_net._all_models)
self._all_params_dict.update(
q_optimizer=self.q_optimizer,
intra_option_optimizer=self.intra_option_optimizer,
termination_optimizer=self.termination_optimizer)
self._model_post_process()
def _generate_random_options(self):
return tf.constant(np.random.randint(0, self.options_num,
self.n_agents),
dtype=tf.int32)
def choose_action(self, s, visual_s, evaluation=False):
if not hasattr(self, 'options'):
self.options = self._generate_random_options()
self.last_options = self.options
a, self.options, self.cell_state = self._get_action(
s, visual_s, self.cell_state, self.options)
if self.use_eps_greedy:
if np.random.uniform() < self.expl_expt_mng.get_esp(
self.train_step, evaluation=evaluation): # epsilon greedy
self.options = self._generate_random_options()
a = a.numpy()
return a
@tf.function
def _get_action(self, s, visual_s, cell_state, options):
with tf.device(self.device):
feat, cell_state = self._representation_net(s,
visual_s,
cell_state=cell_state)
q = self.q_net.value_net(feat) # [B, P]
pi = self.intra_option_net.value_net(feat) # [B, P, A]
beta = self.termination_net.value_net(feat) # [B, P]
options_onehot = tf.one_hot(options,
self.options_num,
dtype=tf.float32) # [B, P]
options_onehot_expanded = tf.expand_dims(options_onehot,
axis=-1) # [B, P, 1]
pi = tf.reduce_sum(pi * options_onehot_expanded, axis=1) # [B, A]
if self.is_continuous:
log_std = tf.gather(self.log_std, options)
mu = tf.math.tanh(pi)
a, _ = gaussian_clip_rsample(mu, log_std)
else:
pi = pi / self.boltzmann_temperature
dist = tfp.distributions.Categorical(
logits=tf.nn.log_softmax(pi)) # [B, ]
a = dist.sample()
max_options = tf.cast(tf.argmax(q, axis=-1),
dtype=tf.int32) # [B, P] => [B, ]
if self.use_eps_greedy:
new_options = max_options
else:
beta_probs = tf.reduce_sum(beta * options_onehot,
axis=1) # [B, P] => [B,]
beta_dist = tfp.distributions.Bernoulli(probs=beta_probs)
new_options = tf.where(beta_dist.sample() < 1, options,
max_options)
return a, new_options, cell_state
def _target_params_update(self):
if self.global_step % self.assign_interval == 0:
update_target_net_weights(self.q_target_net.weights,
self.q_net.weights)
def learn(self, **kwargs):
self.train_step = kwargs.get('train_step')
for i in range(self.train_times_per_step):
self._learn(
function_dict={
'sample_data_list': [
's', 'visual_s', 'a', 'r', 's_', 'visual_s_', 'done',
'last_options', 'options'
],
'train_data_list': [
's', 'visual_s', 'a', 'r', 's_', 'visual_s_', 'done',
'last_options', 'options'
],
'summary_dict':
dict([['LEARNING_RATE/q_lr',
self.q_lr(self.train_step)],
[
'LEARNING_RATE/intra_option_lr',
self.intra_option_lr(self.train_step)
],
[
'LEARNING_RATE/termination_lr',
self.termination_lr(self.train_step)
], ['Statistics/option', self.options[0]]])
})
@tf.function(experimental_relax_shapes=True)
def _train(self, memories, isw, cell_state):
s, visual_s, a, r, s_, visual_s_, done, last_options, options = memories
last_options = tf.cast(last_options, tf.int32)
options = tf.cast(options, tf.int32)
with tf.device(self.device):
with tf.GradientTape(persistent=True) as tape:
feat, _ = self._representation_net(s,
visual_s,
cell_state=cell_state)
feat_, _ = self._representation_target_net(
s_, visual_s_, cell_state=cell_state)
q = self.q_net.value_net(feat) # [B, P]
pi = self.intra_option_net.value_net(feat) # [B, P, A]
beta = self.termination_net.value_net(feat) # [B, P]
q_next = self.q_target_net.value_net(
feat_) # [B, P], [B, P, A], [B, P]
beta_next = self.termination_net.value_net(feat_) # [B, P]
options_onehot = tf.one_hot(options,
self.options_num,
dtype=tf.float32) # [B,] => [B, P]
q_s = qu_eval = tf.reduce_sum(q * options_onehot,
axis=-1,
keepdims=True) # [B, 1]
beta_s_ = tf.reduce_sum(beta_next * options_onehot,
axis=-1,
keepdims=True) # [B, 1]
q_s_ = tf.reduce_sum(q_next * options_onehot,
axis=-1,
keepdims=True) # [B, 1]
# https://github.com/jeanharb/option_critic/blob/5d6c81a650a8f452bc8ad3250f1f211d317fde8c/neural_net.py#L94
if self.double_q:
q_ = self.q_net.value_net(
feat) # [B, P], [B, P, A], [B, P]
max_a_idx = tf.one_hot(
tf.argmax(q_, axis=-1),
self.options_num,
dtype=tf.float32) # [B, P] => [B, ] => [B, P]
q_s_max = tf.reduce_sum(q_next * max_a_idx,
axis=-1,
keepdims=True) # [B, 1]
else:
q_s_max = tf.reduce_max(q_next, axis=-1,
keepdims=True) # [B, 1]
u_target = (1 - beta_s_) * q_s_ + beta_s_ * q_s_max # [B, 1]
qu_target = tf.stop_gradient(r + self.gamma *
(1 - done) * u_target)
td_error = qu_target - qu_eval # gradient : q
q_loss = tf.reduce_mean(tf.square(td_error) *
isw) # [B, 1] => 1
# https://github.com/jeanharb/option_critic/blob/5d6c81a650a8f452bc8ad3250f1f211d317fde8c/neural_net.py#L130
if self.use_baseline:
adv = tf.stop_gradient(qu_target - qu_eval)
else:
adv = tf.stop_gradient(qu_target)
options_onehot_expanded = tf.expand_dims(
options_onehot, axis=-1) # [B, P] => [B, P, 1]
pi = tf.reduce_sum(pi * options_onehot_expanded,
axis=1) # [B, P, A] => [B, A]
if self.is_continuous:
log_std = tf.gather(self.log_std, options)
mu = tf.math.tanh(pi)
log_p = gaussian_likelihood_sum(a, mu, log_std)
entropy = gaussian_entropy(log_std)
else:
pi = pi / self.boltzmann_temperature
log_pi = tf.nn.log_softmax(pi, axis=-1) # [B, A]
entropy = -tf.reduce_sum(tf.exp(log_pi) * log_pi,
axis=1,
keepdims=True) # [B, 1]
log_p = tf.reduce_sum(a * log_pi, axis=-1,
keepdims=True) # [B, 1]
pi_loss = tf.reduce_mean(
-(log_p * adv + self.ent_coff * entropy)
) # [B, 1] * [B, 1] => [B, 1] => 1
last_options_onehot = tf.one_hot(
last_options, self.options_num,
dtype=tf.float32) # [B,] => [B, P]
beta_s = tf.reduce_sum(beta * last_options_onehot,
axis=-1,
keepdims=True) # [B, 1]
if self.use_eps_greedy:
v_s = tf.reduce_max(
q, axis=-1,
keepdims=True) - self.termination_regularizer # [B, 1]
else:
v_s = (1 - beta_s) * q_s + beta_s * tf.reduce_max(
q, axis=-1, keepdims=True) # [B, 1]
# v_s = tf.reduce_mean(q, axis=-1, keepdims=True) # [B, 1]
beta_loss = beta_s * tf.stop_gradient(q_s - v_s) # [B, 1]
# https://github.com/lweitkamp/option-critic-pytorch/blob/0c57da7686f8903ed2d8dded3fae832ee9defd1a/option_critic.py#L238
if self.terminal_mask:
beta_loss *= (1 - done)
beta_loss = tf.reduce_mean(beta_loss) # [B, 1] => 1
q_grads = tape.gradient(q_loss, self.q_net.trainable_variables)
intra_option_grads = tape.gradient(pi_loss, self.actor_tv)
termination_grads = tape.gradient(
beta_loss, self.termination_net.trainable_variables)
self.q_optimizer.apply_gradients(
zip(q_grads, self.q_net.trainable_variables))
self.intra_option_optimizer.apply_gradients(
zip(intra_option_grads, self.actor_tv))
self.termination_optimizer.apply_gradients(
zip(termination_grads,
self.termination_net.trainable_variables))
self.global_step.assign_add(1)
return td_error, dict(
[['LOSS/q_loss', tf.reduce_mean(q_loss)],
['LOSS/pi_loss', tf.reduce_mean(pi_loss)],
['LOSS/beta_loss',
tf.reduce_mean(beta_loss)],
['Statistics/q_option_max',
tf.reduce_max(q_s)],
['Statistics/q_option_min',
tf.reduce_min(q_s)],
['Statistics/q_option_mean',
tf.reduce_mean(q_s)]])
def store_data(self, s, visual_s, a, r, s_, visual_s_, done):
"""
for off-policy training, use this function to store <s, a, r, s_, done> into ReplayBuffer.
"""
assert isinstance(a,
np.ndarray), "store need action type is np.ndarray"
assert isinstance(r,
np.ndarray), "store need reward type is np.ndarray"
assert isinstance(done,
np.ndarray), "store need done type is np.ndarray"
self._running_average(s)
self.data.add(
s,
visual_s,
a,
r[:, np.newaxis], # 升维
s_,
visual_s_,
done[:, np.newaxis], # 升维
self.last_options,
self.options)
def no_op_store(self, s, visual_s, a, r, s_, visual_s_, done):
pass
class TAC(Off_Policy):
"""Tsallis Actor Critic, TAC with V neural Network. https://arxiv.org/abs/1902.00137
"""
def __init__(
self,
envspec,
alpha=0.2,
annealing=True,
last_alpha=0.01,
ployak=0.995,
entropic_index=1.5,
discrete_tau=1.0,
log_std_bound=[-20, 2],
network_settings={
'actor_continuous': {
'share': [128, 128],
'mu': [64],
'log_std': [64]
},
'actor_discrete': [64, 32],
'q': [128, 128]
},
auto_adaption=True,
actor_lr=5.0e-4,
critic_lr=1.0e-3,
alpha_lr=5.0e-4,
**kwargs):
super().__init__(envspec=envspec, **kwargs)
self.ployak = ployak
self.discrete_tau = discrete_tau
self.entropic_index = 2 - entropic_index
self.log_std_min, self.log_std_max = log_std_bound[:]
self.auto_adaption = auto_adaption
self.annealing = annealing
if self.auto_adaption:
self.log_alpha = tf.Variable(initial_value=0.0,
name='log_alpha',
dtype=tf.float32,
trainable=True)
else:
self.log_alpha = tf.Variable(initial_value=tf.math.log(alpha),
name='log_alpha',
dtype=tf.float32,
trainable=False)
if self.annealing:
self.alpha_annealing = LinearAnnealing(alpha, last_alpha, 1e6)
def _create_net(name, representation_net=None):
return DoubleValueNetwork(
name=name,
representation_net=representation_net,
value_net_type=OutputNetworkType.CRITIC_QVALUE_ONE,
value_net_kwargs=dict(action_dim=self.a_dim,
network_settings=network_settings['q']))
self.critic_net = _create_net('critic_net', self._representation_net)
self._representation_target_net = self._create_representation_net(
'_representation_target_net')
self.critic_target_net = _create_net('critic_target_net',
self._representation_target_net)
if self.is_continuous:
self.actor_net = ValueNetwork(
name='actor_net',
value_net_type=OutputNetworkType.ACTOR_CTS,
value_net_kwargs=dict(
vector_dim=self._representation_net.h_dim,
output_shape=self.a_dim,
network_settings=network_settings['actor_continuous']))
else:
self.actor_net = ValueNetwork(
name='actor_net',
value_net_type=OutputNetworkType.ACTOR_DCT,
value_net_kwargs=dict(
vector_dim=self._representation_net.h_dim,
output_shape=self.a_dim,
network_settings=network_settings['actor_discrete']))
self.gumbel_dist = tfp.distributions.Gumbel(0, 1)
# entropy = -log(1/|A|) = log |A|
self.target_entropy = 0.98 * (-self.a_dim if self.is_continuous else
np.log(self.a_dim))
update_target_net_weights(self.critic_target_net.weights,
self.critic_net.weights)
self.actor_lr, self.critic_lr, self.alpha_lr = map(
self.init_lr, [actor_lr, critic_lr, alpha_lr])
self.optimizer_actor, self.optimizer_critic, self.optimizer_alpha = map(
self.init_optimizer,
[self.actor_lr, self.critic_lr, self.alpha_lr])
self._worker_params_dict.update(self._representation_net._all_models)
self._worker_params_dict.update(self.actor_net._policy_models)
self._all_params_dict.update(self.actor_net._all_models)
self._all_params_dict.update(self.critic_net._all_models)
self._all_params_dict.update(log_alpha=self.log_alpha,
optimizer_actor=self.optimizer_actor,
optimizer_critic=self.optimizer_critic,
optimizer_alpha=self.optimizer_alpha)
self._model_post_process()
@property
def alpha(self):
return tf.exp(self.log_alpha)
def choose_action(self, s, visual_s, evaluation=False):
mu, pi, self.cell_state = self._get_action(s, visual_s,
self.cell_state)
a = mu.numpy() if evaluation else pi.numpy()
return a
@tf.function
def _get_action(self, s, visual_s, cell_state):
with tf.device(self.device):
feat, cell_state = self._representation_net(s,
visual_s,
cell_state=cell_state)
if self.is_continuous:
mu, log_std = self.actor_net.value_net(feat)
log_std = clip_nn_log_std(log_std, self.log_std_min,
self.log_std_max)
pi, _ = tsallis_squash_rsample(mu, log_std,
self.entropic_index)
mu = tf.tanh(mu) # squash mu
else:
logits = self.actor_net.value_net(feat)
mu = tf.argmax(logits, axis=1)
cate_dist = tfp.distributions.Categorical(
logits=tf.nn.log_softmax(logits))
pi = cate_dist.sample()
return mu, pi, cell_state
def _target_params_update(self):
update_target_net_weights(self.critic_target_net.weights,
self.critic_net.weights, self.ployak)
def learn(self, **kwargs):
self.train_step = kwargs.get('train_step')
for i in range(self.train_times_per_step):
self._learn(
function_dict={
'summary_dict':
dict([[
'LEARNING_RATE/actor_lr',
self.actor_lr(self.train_step)
],
[
'LEARNING_RATE/critic_lr',
self.critic_lr(self.train_step)
],
[
'LEARNING_RATE/alpha_lr',
self.alpha_lr(self.train_step)
]]),
'train_data_list':
['ss', 'vvss', 'a', 'r', 'done', 's_', 'visual_s_']
})
def _train(self, memories, isw, cell_state):
td_error, summaries = self.train(memories, isw, cell_state)
if self.annealing and not self.auto_adaption:
self.log_alpha.assign(
tf.math.log(
tf.cast(self.alpha_annealing(self.global_step.numpy()),
tf.float32)))
return td_error, summaries
@tf.function(experimental_relax_shapes=True)
def train(self, memories, isw, cell_state):
ss, vvss, a, r, done, s_, visual_s_ = memories
with tf.device(self.device):
with tf.GradientTape(persistent=True) as tape:
(feat,
feat_), _ = self._representation_net(ss,
vvss,
cell_state=cell_state,
need_split=True)
if self.is_continuous:
mu, log_std = self.actor_net.value_net(feat)
log_std = clip_nn_log_std(log_std, self.log_std_min,
self.log_std_max)
pi, log_pi = tsallis_squash_rsample(
mu, log_std, self.entropic_index)
entropy = gaussian_entropy(log_std)
target_mu, target_log_std = self.actor_net.value_net(feat_)
target_log_std = clip_nn_log_std(target_log_std,
self.log_std_min,
self.log_std_max)
target_pi, target_log_pi = tsallis_squash_rsample(
target_mu, target_log_std, self.entropic_index)
else:
logits = self.actor_net.value_net(feat)
logp_all = tf.nn.log_softmax(logits)
gumbel_noise = tf.cast(self.gumbel_dist.sample(a.shape),
dtype=tf.float32)
_pi = tf.nn.softmax(
(logp_all + gumbel_noise) / self.discrete_tau)
_pi_true_one_hot = tf.one_hot(tf.argmax(_pi, axis=-1),
self.a_dim)
_pi_diff = tf.stop_gradient(_pi_true_one_hot - _pi)
pi = _pi_diff + _pi
log_pi = tf.reduce_sum(tf.multiply(logp_all, pi),
axis=1,
keepdims=True)
entropy = -tf.reduce_mean(
tf.reduce_sum(tf.exp(logp_all) * logp_all,
axis=1,
keepdims=True))
target_logits = self.actor_net.value_net(feat_)
target_cate_dist = tfp.distributions.Categorical(
logits=tf.nn.log_softmax(target_logits))
target_pi = target_cate_dist.sample()
target_log_pi = target_cate_dist.log_prob(target_pi)
target_pi = tf.one_hot(target_pi,
self.a_dim,
dtype=tf.float32)
q1, q2 = self.critic_net.get_value(feat, a)
q_s_pi = self.critic_net.get_min(feat, pi)
q1_target, q2_target, _ = self.critic_target_net(
s_, visual_s_, target_pi, cell_state=cell_state)
q_target = tf.minimum(q1_target, q2_target)
dc_r = tf.stop_gradient(
r + self.gamma * (1 - done) *
(q_target - self.alpha * target_log_pi))
td_error1 = q1 - dc_r
td_error2 = q2 - dc_r
q1_loss = tf.reduce_mean(tf.square(td_error1) * isw)
q2_loss = tf.reduce_mean(tf.square(td_error2) * isw)
critic_loss = 0.5 * q1_loss + 0.5 * q2_loss
actor_loss = -tf.reduce_mean(q_s_pi - self.alpha * log_pi)
if self.auto_adaption:
alpha_loss = -tf.reduce_mean(
self.alpha *
tf.stop_gradient(log_pi + self.target_entropy))
critic_grads = tape.gradient(critic_loss,
self.critic_net.trainable_variables)
self.optimizer_critic.apply_gradients(
zip(critic_grads, self.critic_net.trainable_variables))
actor_grads = tape.gradient(actor_loss,
self.actor_net.trainable_variables)
self.optimizer_actor.apply_gradients(
zip(actor_grads, self.actor_net.trainable_variables))
if self.auto_adaption:
alpha_grad = tape.gradient(alpha_loss, self.log_alpha)
self.optimizer_alpha.apply_gradients([(alpha_grad,
self.log_alpha)])
self.global_step.assign_add(1)
summaries = dict(
[['LOSS/actor_loss', actor_loss], ['LOSS/q1_loss', q1_loss],
['LOSS/q2_loss', q2_loss], ['LOSS/critic_loss', critic_loss],
['Statistics/log_alpha', self.log_alpha],
['Statistics/alpha', self.alpha],
['Statistics/entropy', entropy],
['Statistics/q_min',
tf.reduce_min(tf.minimum(q1, q2))],
['Statistics/q_mean',
tf.reduce_mean(tf.minimum(q1, q2))],
['Statistics/q_max',
tf.reduce_max(tf.maximum(q1, q2))]])
if self.auto_adaption:
summaries.update({'LOSS/alpha_loss': alpha_loss})
return (td_error1 + td_error2) / 2, summaries
class SAC_V(Off_Policy):
"""
Soft Actor Critic with Value neural network. https://arxiv.org/abs/1812.05905
Soft Actor-Critic for Discrete Action Settings. https://arxiv.org/abs/1910.07207
"""
def __init__(
self,
envspec,
alpha=0.2,
annealing=True,
last_alpha=0.01,
ployak=0.995,
use_gumbel=True,
discrete_tau=1.0,
network_settings={
'actor_continuous': {
'share': [128, 128],
'mu': [64],
'log_std': [64],
'soft_clip': False,
'log_std_bound': [-20, 2]
},
'actor_discrete': [64, 32],
'q': [128, 128],
'v': [128, 128]
},
actor_lr=5.0e-4,
critic_lr=1.0e-3,
alpha_lr=5.0e-4,
auto_adaption=True,
**kwargs):
super().__init__(envspec=envspec, **kwargs)
self.ployak = ployak
self.use_gumbel = use_gumbel
self.discrete_tau = discrete_tau
self.auto_adaption = auto_adaption
self.annealing = annealing
if self.auto_adaption:
self.log_alpha = tf.Variable(initial_value=0.0,
name='log_alpha',
dtype=tf.float32,
trainable=True)
else:
self.log_alpha = tf.Variable(initial_value=tf.math.log(alpha),
name='log_alpha',
dtype=tf.float32,
trainable=False)
if self.annealing:
self.alpha_annealing = LinearAnnealing(alpha, last_alpha, 1e6)
def _create_net(name, representation_net=None):
return ValueNetwork(
name=name,
representation_net=representation_net,
value_net_type=OutputNetworkType.CRITIC_VALUE,
value_net_kwargs=dict(network_settings=network_settings['v']))
self.v_net = _create_net('v_net', self._representation_net)
self._representation_target_net = self._create_representation_net(
'_representation_target_net')
self.v_target_net = _create_net('v_target_net',
self._representation_target_net)
if self.is_continuous:
self.actor_net = ValueNetwork(
name='actor_net',
value_net_type=OutputNetworkType.ACTOR_CTS,
value_net_kwargs=dict(
vector_dim=self._representation_net.h_dim,
output_shape=self.a_dim,
network_settings=network_settings['actor_continuous']))
else:
self.actor_net = ValueNetwork(
name='actor_net',
value_net_type=OutputNetworkType.ACTOR_DCT,
value_net_kwargs=dict(
vector_dim=self._representation_net.h_dim,
output_shape=self.a_dim,
network_settings=network_settings['actor_discrete']))
if self.use_gumbel:
self.gumbel_dist = tfp.distributions.Gumbel(0, 1)
# entropy = -log(1/|A|) = log |A|
self.target_entropy = 0.98 * (-self.a_dim if self.is_continuous else
np.log(self.a_dim))
if self.is_continuous or self.use_gumbel:
self.q_net = DoubleValueNetwork(
name='q_net',
value_net_type=OutputNetworkType.CRITIC_QVALUE_ONE,
value_net_kwargs=dict(
vector_dim=self._representation_net.h_dim,
action_dim=self.a_dim,
network_settings=network_settings['q']))
else:
self.q_net = DoubleValueNetwork(
name='q_net',
value_net_type=OutputNetworkType.CRITIC_QVALUE_ALL,
value_net_kwargs=dict(
vector_dim=self._representation_net.h_dim,
output_shape=self.a_dim,
network_settings=network_settings['q']))
update_target_net_weights(self.v_target_net.weights,
self.v_net.weights)
self.actor_lr, self.critic_lr, self.alpha_lr = map(
self.init_lr, [actor_lr, critic_lr, alpha_lr])
self.optimizer_actor, self.optimizer_critic, self.optimizer_alpha = map(
self.init_optimizer,
[self.actor_lr, self.critic_lr, self.alpha_lr])
self._worker_params_dict.update(self._representation_net._all_models)
self._worker_params_dict.update(self.actor_net._policy_models)
self._all_params_dict.update(self.actor_net._all_models)
self._all_params_dict.update(self.v_net._all_models)
self._all_params_dict.update(self.q_net._all_models)
self._all_params_dict.update(log_alpha=self.log_alpha,
optimizer_actor=self.optimizer_actor,
optimizer_critic=self.optimizer_critic,
optimizer_alpha=self.optimizer_alpha)
self._model_post_process()
@property
def alpha(self):
return tf.exp(self.log_alpha)
def choose_action(self, obs, evaluation=False):
mu, pi, self.cell_state = self._get_action(obs, self.cell_state)
a = mu.numpy() if evaluation else pi.numpy()
return a
@tf.function
def _get_action(self, obs, cell_state):
with tf.device(self.device):
feat, cell_state = self._representation_net(obs,
cell_state=cell_state)
if self.is_continuous:
mu, log_std = self.actor_net.value_net(feat)
pi, _ = squash_rsample(mu, log_std)
mu = tf.tanh(mu) # squash mu
else:
logits = self.actor_net.value_net(feat)
mu = tf.argmax(logits, axis=1)
cate_dist = tfp.distributions.Categorical(
logits=tf.nn.log_softmax(logits))
pi = cate_dist.sample()
return mu, pi, cell_state
def _target_params_update(self):
update_target_net_weights(self.v_target_net.weights,
self.v_net.weights, self.ployak)
def learn(self, **kwargs):
self.train_step = kwargs.get('train_step')
for i in range(self.train_times_per_step):
self._learn(
function_dict={
'summary_dict':
dict([[
'LEARNING_RATE/actor_lr',
self.actor_lr(self.train_step)
],
[
'LEARNING_RATE/critic_lr',
self.critic_lr(self.train_step)
],
[
'LEARNING_RATE/alpha_lr',
self.alpha_lr(self.train_step)
]]),
})
def _train(self, BATCH, isw, cell_state):
if self.is_continuous or self.use_gumbel:
td_error, summaries = self.train_continuous(BATCH, isw, cell_state)
else:
td_error, summaries = self.train_discrete(BATCH, isw, cell_state)
if self.annealing and not self.auto_adaption:
self.log_alpha.assign(
tf.math.log(
tf.cast(self.alpha_annealing(self.global_step.numpy()),
tf.float32)))
return td_error, summaries
@tf.function
def train_continuous(self, BATCH, isw, cell_state):
with tf.device(self.device):
with tf.GradientTape(persistent=True) as tape:
feat, _ = self._representation_net(BATCH.obs,
cell_state=cell_state)
v = self.v_net.value_net(feat)
v_target, _ = self.v_target_net(BATCH.obs_,
cell_state=cell_state)
if self.is_continuous:
mu, log_std = self.actor_net.value_net(feat)
pi, log_pi = squash_rsample(mu, log_std)
entropy = gaussian_entropy(log_std)
else:
logits = self.actor_net.value_net(feat)
logp_all = tf.nn.log_softmax(logits)
gumbel_noise = tf.cast(self.gumbel_dist.sample(
BATCH.action.shape),
dtype=tf.float32)
_pi = tf.nn.softmax(
(logp_all + gumbel_noise) / self.discrete_tau)
_pi_true_one_hot = tf.one_hot(tf.argmax(_pi, axis=-1),
self.a_dim)
_pi_diff = tf.stop_gradient(_pi_true_one_hot - _pi)
pi = _pi_diff + _pi
log_pi = tf.reduce_sum(tf.multiply(logp_all, pi),
axis=1,
keepdims=True)
entropy = -tf.reduce_mean(
tf.reduce_sum(tf.exp(logp_all) * logp_all,
axis=1,
keepdims=True))
q1, q2 = self.q_net.get_value(feat, BATCH.action)
q1_pi, q2_pi = self.q_net.get_value(feat, pi)
dc_r = tf.stop_gradient(BATCH.reward + self.gamma * v_target *
(1 - BATCH.done))
v_from_q_stop = tf.stop_gradient(
tf.minimum(q1_pi, q2_pi) - self.alpha * log_pi)
td_v = v - v_from_q_stop
td_error1 = q1 - dc_r
td_error2 = q2 - dc_r
q1_loss = tf.reduce_mean(tf.square(td_error1) * isw)
q2_loss = tf.reduce_mean(tf.square(td_error2) * isw)
v_loss_stop = tf.reduce_mean(tf.square(td_v) * isw)
critic_loss = 0.5 * q1_loss + 0.5 * q2_loss + 0.5 * v_loss_stop
actor_loss = -tf.reduce_mean(q1_pi - self.alpha * log_pi)
if self.auto_adaption:
alpha_loss = -tf.reduce_mean(
self.alpha *
tf.stop_gradient(log_pi + self.target_entropy))
actor_grads = tape.gradient(actor_loss,
self.actor_net.trainable_variables)
self.optimizer_actor.apply_gradients(
zip(actor_grads, self.actor_net.trainable_variables))
critic_grads = tape.gradient(
critic_loss, self.q_net.trainable_variables +
self.v_net.trainable_variables)
self.optimizer_critic.apply_gradients(
zip(
critic_grads, self.q_net.trainable_variables +
self.v_net.trainable_variables))
if self.auto_adaption:
alpha_grad = tape.gradient(alpha_loss, self.log_alpha)
self.optimizer_alpha.apply_gradients([(alpha_grad,
self.log_alpha)])
self.global_step.assign_add(1)
summaries = dict(
[['LOSS/actor_loss', actor_loss], ['LOSS/q1_loss', q1_loss],
['LOSS/q2_loss', q2_loss], ['LOSS/v_loss', v_loss_stop],
['LOSS/critic_loss', critic_loss],
['Statistics/log_alpha', self.log_alpha],
['Statistics/alpha', self.alpha],
['Statistics/entropy', entropy],
['Statistics/q_min',
tf.reduce_min(tf.minimum(q1, q2))],
['Statistics/q_mean',
tf.reduce_mean(tf.minimum(q1, q2))],
['Statistics/q_max',
tf.reduce_max(tf.maximum(q1, q2))],
['Statistics/v_mean', tf.reduce_mean(v)]])
if self.auto_adaption:
summaries.update({'LOSS/alpha_loss': alpha_loss})
return (td_error1 + td_error2) / 2, summaries
@tf.function
def train_discrete(self, BATCH, isw, cell_state):
with tf.device(self.device):
with tf.GradientTape(persistent=True) as tape:
feat, _ = self._representation_net(BATCH.obs,
cell_state=cell_state)
v = self.v_net.value_net(feat) # [B, 1]
v_target, _ = self.v_target_net(
BATCH.obs_, cell_state=cell_state) # [B, 1]
q1_all, q2_all = self.q_net.get_value(feat) # [B, A]
def q_function(x):
return tf.reduce_sum(x * BATCH.action,
axis=-1,
keepdims=True) # [B, 1]
q1 = q_function(q1_all)
q2 = q_function(q2_all)
logits = self.actor_net.value_net(feat) # [B, A]
logp_all = tf.nn.log_softmax(logits) # [B, A]
entropy = -tf.reduce_sum(tf.exp(logp_all) * logp_all,
axis=1,
keepdims=True) # [B, 1]
q_all = self.q_net.get_min(feat) # [B, A]
actor_loss = -tf.reduce_mean(
tf.reduce_sum((q_all - self.alpha * logp_all) *
tf.exp(logp_all)) # [B, A] => [B,]
)
dc_r = tf.stop_gradient(BATCH.reward + self.gamma * v_target *
(1 - BATCH.done))
td_v = v - tf.stop_gradient(
tf.minimum(
tf.reduce_sum(tf.exp(logp_all) * q1_all, axis=-1),
tf.reduce_sum(tf.exp(logp_all) * q2_all, axis=-1)))
td_error1 = q1 - dc_r
td_error2 = q2 - dc_r
q1_loss = tf.reduce_mean(tf.square(td_error1) * isw)
q2_loss = tf.reduce_mean(tf.square(td_error2) * isw)
v_loss_stop = tf.reduce_mean(tf.square(td_v) * isw)
critic_loss = 0.5 * q1_loss + 0.5 * q2_loss + 0.5 * v_loss_stop
if self.auto_adaption:
corr = tf.stop_gradient(self.target_entropy - entropy)
# corr = tf.stop_gradient(tf.reduce_sum((logp_all - self.a_dim) * tf.exp(logp_all), axis=-1)) #[B, A] => [B,]
alpha_loss = -tf.reduce_mean(self.alpha * corr)
critic_grads = tape.gradient(
critic_loss, self.q_net.trainable_variables +
self.v_net.trainable_variables)
self.optimizer_critic.apply_gradients(
zip(
critic_grads, self.q_net.trainable_variables +
self.v_net.trainable_variables))
actor_grads = tape.gradient(actor_loss,
self.actor_net.trainable_variables)
self.optimizer_actor.apply_gradients(
zip(actor_grads, self.actor_net.trainable_variables))
if self.auto_adaption:
alpha_grad = tape.gradient(alpha_loss, self.log_alpha)
self.optimizer_alpha.apply_gradients([(alpha_grad,
self.log_alpha)])
self.global_step.assign_add(1)
summaries = dict([['LOSS/actor_loss', actor_loss],
['LOSS/q1_loss', q1_loss],
['LOSS/q2_loss', q2_loss],
['LOSS/v_loss', v_loss_stop],
['LOSS/critic_loss', critic_loss],
['Statistics/log_alpha', self.log_alpha],
['Statistics/alpha', self.alpha],
['Statistics/entropy',
tf.reduce_mean(entropy)],
['Statistics/v_mean',
tf.reduce_mean(v)]])
if self.auto_adaption:
summaries.update({'LOSS/alpha_loss': alpha_loss})
return (td_error1 + td_error2) / 2, summaries