def __init__(
self,
logits: tf.Tensor,
act_size: List[int],
reparameterize: bool = False,
tanh_squash: bool = False,
condition_sigma: bool = True,
log_sigma_min: float = -20,
log_sigma_max: float = 2,
):
"""
A Gaussian output distribution for continuous actions.
:param logits: Hidden layer to use as the input to the Gaussian distribution.
:param act_size: List containing the number of continuous actions.
:param reparameterize: Whether or not to use the reparameterization trick (block gradients through
log probability calculation.)
:param tanh_squash: Squash the output using tanh, constraining it between -1 and 1.
From: Haarnoja et. al, https://arxiv.org/abs/1801.01290
:param log_sigma_min: Minimum log standard deviation to clip by.
:param log_sigma_max: Maximum log standard deviation to clip by.
"""
encoded = self._create_mu_log_sigma(
logits,
act_size,
log_sigma_min,
log_sigma_max,
condition_sigma=condition_sigma,
)
self._sampled_policy = self._create_sampled_policy(encoded)
if not reparameterize:
_sampled_policy_probs = tf.stop_gradient(self._sampled_policy)
else:
_sampled_policy_probs = self._sampled_policy
self._all_probs = self._create_log_probs(_sampled_policy_probs,
encoded)
if tanh_squash:
self._sampled_policy = tf.tanh(self._sampled_policy)
self._all_probs = self._do_squash_correction_for_tanh(
self._all_probs, self._sampled_policy)
self._total_prob = tf.reduce_sum(self._all_probs,
axis=1,
keepdims=True)
self._entropy = self._create_entropy(encoded)
def create_cc_actor(self, hidden_policy, scope):
"""
Creates Continuous control actor for SAC.
:param hidden_policy: Output of feature extractor (i.e. the input for vector obs, output of CNN for visual obs).
:param num_layers: TF scope to assign whatever is created in this block.
"""
# Create action input (continuous)
self.action_holder = tf.placeholder(shape=[None, self.act_size[0]],
dtype=tf.float32,
name="action_holder")
self.external_action_in = self.action_holder
scope = self.join_scopes(scope, "policy")
with tf.variable_scope(scope):
hidden_policy = self.create_vector_observation_encoder(
hidden_policy,
self.h_size,
self.activ_fn,
self.num_layers,
"encoder",
False,
)
if self.use_recurrent:
hidden_policy, memory_out = self.create_recurrent_encoder(
hidden_policy,
self.policy_memory_in,
self.sequence_length,
name="lstm_policy",
)
self.policy_memory_out = memory_out
with tf.variable_scope(scope):
mu = tf.layers.dense(
hidden_policy,
self.act_size[0],
activation=None,
name="mu",
kernel_initializer=LearningModel.scaled_init(0.01),
)
# Policy-dependent log_sigma_sq
log_sigma_sq = tf.layers.dense(
hidden_policy,
self.act_size[0],
activation=None,
name="log_std",
kernel_initializer=LearningModel.scaled_init(0.01),
)
self.log_sigma_sq = tf.clip_by_value(log_sigma_sq, LOG_STD_MIN,
LOG_STD_MAX)
sigma_sq = tf.exp(self.log_sigma_sq)
# Do the reparameterization trick
policy_ = mu + tf.random_normal(tf.shape(mu)) * sigma_sq
_gauss_pre = -0.5 * (((policy_ - mu) /
(tf.exp(self.log_sigma_sq) + EPSILON))**2 +
2 * self.log_sigma_sq + np.log(2 * np.pi))
all_probs = tf.reduce_sum(_gauss_pre, axis=1, keepdims=True)
self.entropy = tf.reduce_sum(self.log_sigma_sq +
0.5 * np.log(2.0 * np.pi * np.e),
axis=-1)
# Squash probabilities
# Keep deterministic around in case we want to use it.
self.deterministic_output = tf.tanh(mu)
# Note that this is just for symmetry with PPO.
self.output_pre = tf.tanh(policy_)
# Squash correction
all_probs -= tf.reduce_sum(tf.log(1 - self.output_pre**2 +
EPSILON),
axis=1,
keepdims=True)
self.all_log_probs = all_probs
self.selected_actions = tf.stop_gradient(self.output_pre)
self.action_probs = all_probs
# Extract output for Barracuda
self.output = tf.identity(self.output_pre, name="action")
# Get all policy vars
self.policy_vars = self.get_vars(scope)
def _create_cc_actor(
self,
encoded: tf.Tensor,
tanh_squash: bool = False,
reparameterize: bool = False,
condition_sigma_on_obs: bool = True,
) -> None:
"""
Creates Continuous control actor-critic model.
:param h_size: Size of hidden linear layers.
:param num_layers: Number of hidden linear layers.
:param vis_encode_type: Type of visual encoder to use if visual input.
:param tanh_squash: Whether to use a tanh function, or a clipped output.
:param reparameterize: Whether we are using the resampling trick to update the policy.
"""
if self.use_recurrent:
self.memory_in = tf.placeholder(shape=[None, self.m_size],
dtype=tf.float32,
name="recurrent_in")
hidden_policy, memory_policy_out = ModelUtils.create_recurrent_encoder(
encoded,
self.memory_in,
self.sequence_length_ph,
name="lstm_policy")
self.memory_out = tf.identity(memory_policy_out,
name="recurrent_out")
else:
hidden_policy = encoded
with tf.variable_scope("policy"):
mu = tf.layers.dense(
hidden_policy,
self.act_size[0],
activation=None,
name="mu",
kernel_initializer=ModelUtils.scaled_init(0.01),
reuse=tf.AUTO_REUSE,
)
# Policy-dependent log_sigma
if condition_sigma_on_obs:
log_sigma = tf.layers.dense(
hidden_policy,
self.act_size[0],
activation=None,
name="log_sigma",
kernel_initializer=ModelUtils.scaled_init(0.01),
)
else:
log_sigma = tf.get_variable(
"log_sigma",
[self.act_size[0]],
dtype=tf.float32,
initializer=tf.zeros_initializer(),
)
log_sigma = tf.clip_by_value(log_sigma, self.log_std_min,
self.log_std_max)
sigma = tf.exp(log_sigma)
epsilon = tf.random_normal(tf.shape(mu))
sampled_policy = mu + sigma * epsilon
# Stop gradient if we're not doing the resampling trick
if not reparameterize:
sampled_policy_probs = tf.stop_gradient(sampled_policy)
else:
sampled_policy_probs = sampled_policy
# Compute probability of model output.
_gauss_pre = -0.5 * (
((sampled_policy_probs - mu) /
(sigma + EPSILON))**2 + 2 * log_sigma + np.log(2 * np.pi))
all_probs = _gauss_pre
all_probs = tf.reduce_sum(_gauss_pre, axis=1, keepdims=True)
if tanh_squash:
self.output_pre = tf.tanh(sampled_policy)
# Squash correction
all_probs -= tf.reduce_sum(tf.log(1 - self.output_pre**2 +
EPSILON),
axis=1,
keepdims=True)
self.output = tf.identity(self.output_pre, name="action")
else:
self.output_pre = sampled_policy
# Clip and scale output to ensure actions are always within [-1, 1] range.
output_post = tf.clip_by_value(self.output_pre, -3, 3) / 3
self.output = tf.identity(output_post, name="action")
self.selected_actions = tf.stop_gradient(self.output)
self.all_log_probs = tf.identity(all_probs, name="action_probs")
single_dim_entropy = 0.5 * tf.reduce_mean(
tf.log(2 * np.pi * np.e) + 2 * log_sigma)
# Make entropy the right shape
self.entropy = tf.ones_like(tf.reshape(mu[:, 0],
[-1])) * single_dim_entropy
# We keep these tensors the same name, but use new nodes to keep code parallelism with discrete control.
self.log_probs = tf.reduce_sum((tf.identity(self.all_log_probs)),
axis=1,
keepdims=True)
self.action_holder = tf.placeholder(shape=[None, self.act_size[0]],
dtype=tf.float32,
name="action_holder")
