wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-07-09 07:59:29 +08:00
Code Issues Packages Projects Releases Wiki Activity

[RLlib] SAC refactor with new SquashedGaussian distribution class. (#7272)

Browse Source
This commit is contained in:
Sven Mika
2020-02-24 01:10:20 +01:00
committed by GitHub
parent 1660b52751
commit e1fc8368d4
7 changed files with 38 additions and 147 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rllib/BUILD
+1 -1
View File
@@ -918,7 +918,7 @@ py_test(
py_test(
name = "tests/test_explorations",
tags = ["tests_dir", "tests_dir_E", "explorations"],
size = "medium",
size = "large",
srcs = ["tests/test_explorations.py"]
)
rllib/agents/sac/__init__.py
+2 -3
View File
@@ -1,9 +1,8 @@
from ray.rllib.agents.sac.sac import SACTrainer, DEFAULT_CONFIG
from ray.rllib.utils import renamed_agent
SACAgent = renamed_agent(SACTrainer)
from ray.rllib.agents.sac.sac_policy import SACTFPolicy
__all__ = [
"SACTFPolicy",
"SACTrainer",
"DEFAULT_CONFIG",
]
rllib/agents/sac/sac.py
+1 -1
View File
@@ -25,7 +25,7 @@ DEFAULT_CONFIG = with_common_config({
"hidden_activation": "relu",
"hidden_layer_sizes": (256, 256),
},
# Unsquash actions to the upper and lower bounds of env's action space
# Unsquash actions to the upper and lower bounds of env's action space.
"normalize_actions": True,
# === Learning ===
rllib/agents/sac/sac_model.py
+8 -102
View File
@@ -75,96 +75,23 @@ class SACModel(TFModelV2):
self.action_dim = np.product(action_space.shape)
self.model_out = tf.keras.layers.Input(
shape=(num_outputs, ), name="model_out")
self.actions = tf.keras.layers.Input(
shape=(self.action_dim, ), name="actions")
shift_and_log_scale_diag = tf.keras.Sequential([
self.action_model = tf.keras.Sequential([
tf.keras.layers.Dense(
units=hidden,
activation=getattr(tf.nn, actor_hidden_activation, None),
name="action_hidden_{}".format(i))
name="action_{}".format(i + 1))
for i, hidden in enumerate(actor_hiddens)
] + [
tf.keras.layers.Dense(
units=2 * self.action_dim, activation=None, name="action_out")
])(self.model_out)
])
self.shift_and_log_scale_diag = self.action_model(self.model_out)
shift, log_scale_diag = tf.keras.layers.Lambda(
lambda shift_and_log_scale_diag: tf.split(
shift_and_log_scale_diag,
num_or_size_splits=2,
axis=-1)
)(shift_and_log_scale_diag)
log_scale_diag = tf.keras.layers.Lambda(
lambda log_sd: tf.clip_by_value(log_sd, *SCALE_DIAG_MIN_MAX))(
log_scale_diag)
shift_and_log_scale_diag = tf.keras.layers.Concatenate(axis=-1)(
[shift, log_scale_diag])
batch_size = tf.keras.layers.Lambda(lambda x: tf.shape(input=x)[0])(
self.model_out)
base_distribution = tfp.distributions.MultivariateNormalDiag(
loc=tf.zeros(self.action_dim), scale_diag=tf.ones(self.action_dim))
latents = tf.keras.layers.Lambda(
lambda batch_size: base_distribution.sample(batch_size))(
batch_size)
self.shift_and_log_scale_diag = latents
self.latents_model = tf.keras.Model(self.model_out, latents)
def raw_actions_fn(inputs):
shift, log_scale_diag, latents = inputs
bijector = tfp.bijectors.Affine(
shift=shift, scale_diag=tf.exp(log_scale_diag))
actions = bijector.forward(latents)
return actions
raw_actions = tf.keras.layers.Lambda(raw_actions_fn)(
(shift, log_scale_diag, latents))
squash_bijector = (SquashBijector())
actions = tf.keras.layers.Lambda(
lambda raw_actions: squash_bijector.forward(raw_actions))(
raw_actions)
self.actions_model = tf.keras.Model(self.model_out, actions)
deterministic_actions = tf.keras.layers.Lambda(
lambda shift: squash_bijector.forward(shift))(shift)
self.deterministic_actions_model = tf.keras.Model(
self.model_out, deterministic_actions)
def log_pis_fn(inputs):
shift, log_scale_diag, actions = inputs
base_distribution = tfp.distributions.MultivariateNormalDiag(
loc=tf.zeros(self.action_dim),
scale_diag=tf.ones(self.action_dim))
bijector = tfp.bijectors.Chain((
squash_bijector,
tfp.bijectors.Affine(
shift=shift, scale_diag=tf.exp(log_scale_diag)),
))
distribution = (tfp.distributions.TransformedDistribution(
distribution=base_distribution, bijector=bijector))
log_pis = distribution.log_prob(actions)[:, None]
return log_pis
self.register_variables(self.action_model.variables)
self.actions_input = tf.keras.layers.Input(
shape=(self.action_dim, ), name="actions")
log_pis_for_action_input = tf.keras.layers.Lambda(log_pis_fn)(
[shift, log_scale_diag, self.actions_input])
self.log_pis_model = tf.keras.Model(
(self.model_out, self.actions_input), log_pis_for_action_input)
self.register_variables(self.actions_model.variables)
def build_q_net(name, observations, actions):
q_net = tf.keras.Sequential([
tf.keras.layers.Concatenate(axis=1),
@@ -184,12 +111,12 @@ class SACModel(TFModelV2):
q_net([observations, actions]))
return q_net
self.q_net = build_q_net("q", self.model_out, self.actions)
self.q_net = build_q_net("q", self.model_out, self.actions_input)
self.register_variables(self.q_net.variables)
if twin_q:
self.twin_q_net = build_q_net("twin_q", self.model_out,
self.actions)
self.actions_input)
self.register_variables(self.twin_q_net.variables)
else:
self.twin_q_net = None
@@ -199,27 +126,6 @@ class SACModel(TFModelV2):
self.register_variables([self.log_alpha])
def get_policy_output(self, model_out, deterministic=False):
"""Return the (unscaled) output of the policy network.
This returns the unscaled outputs of pi(s).
Arguments:
model_out (Tensor): obs embeddings from the model layers, of shape
[BATCH_SIZE, num_outputs].
Returns:
tensor of shape [BATCH_SIZE, action_dim] with range [-inf, inf].
"""
if deterministic:
actions = self.deterministic_actions_model(model_out)
log_pis = None
else:
actions = self.actions_model(model_out)
log_pis = self.log_pis_model((model_out, actions))
return actions, log_pis
def get_q_values(self, model_out, actions):
"""Return the Q estimates for the most recent forward pass.
@@ -257,7 +163,7 @@ class SACModel(TFModelV2):
def policy_variables(self):
"""Return the list of variables for the policy net."""
return list(self.actions_model.variables)
return list(self.action_model.variables)
def q_variables(self):
"""Return the list of variables for Q / twin Q nets."""
rllib/agents/sac/sac_policy.py
+26 -38
View File
@@ -12,6 +12,7 @@ from ray.rllib.policy.sample_batch import SampleBatch
from ray.rllib.policy.tf_policy import TFPolicy
from ray.rllib.policy.tf_policy_template import build_tf_policy
from ray.rllib.models import ModelCatalog
from ray.rllib.models.tf.tf_action_dist import SquashedGaussian, DiagGaussian
from ray.rllib.utils.error import UnsupportedSpaceException
from ray.rllib.utils import try_import_tf, try_import_tfp
from ray.rllib.utils.annotations import override
@@ -19,6 +20,7 @@ from ray.rllib.utils.tf_ops import minimize_and_clip, make_tf_callable
tf = try_import_tf()
tfp = try_import_tfp()
logger = logging.getLogger(__name__)
@@ -86,18 +88,10 @@ def postprocess_trajectory(policy,
return postprocess_nstep_and_prio(policy, sample_batch)
def unsquash_actions(actions, action_space):
# Use sigmoid to scale to [0,1], but also double magnitude of input to
# emulate behaviour of tanh activation used in SAC and TD3 papers.
sigmoid_out = tf.nn.sigmoid(2 * actions)
# Rescale to actual env policy scale
# (shape of sigmoid_out is [batch_size, dim_actions], so we reshape to
# get same dims)
action_range = (action_space.high - action_space.low)[None]
low_action = action_space.low[None]
unsquashed_actions = action_range * sigmoid_out + low_action
return unsquashed_actions
def get_dist_class(config, action_space):
action_dist_class = SquashedGaussian if \
config["normalize_actions"] is True else DiagGaussian
return action_dist_class
def get_log_likelihood(policy, model, actions, input_dict, obs_space,
@@ -106,8 +100,9 @@ def get_log_likelihood(policy, model, actions, input_dict, obs_space,
"obs": input_dict[SampleBatch.CUR_OBS],
"is_training": policy._get_is_training_placeholder(),
}, [], None)
log_pis = policy.model.log_pis_model((model_out, actions))
return log_pis
distribution_inputs = model.action_model(model_out)
action_dist_class = get_dist_class(policy.config, action_space)
return action_dist_class(distribution_inputs, model).logp(actions)
def build_action_output(policy, model, input_dict, obs_space, action_space,
@@ -116,29 +111,14 @@ def build_action_output(policy, model, input_dict, obs_space, action_space,
"obs": input_dict[SampleBatch.CUR_OBS],
"is_training": policy._get_is_training_placeholder(),
}, [], None)
distribution_inputs = model.action_model(model_out)
action_dist_class = get_dist_class(policy.config, action_space)
squashed_stochastic_actions, log_pis = policy.model.get_policy_output(
model_out, deterministic=False)
stochastic_actions = squashed_stochastic_actions if config[
"normalize_actions"] else unsquash_actions(squashed_stochastic_actions,
action_space)
squashed_deterministic_actions, _ = policy.model.get_policy_output(
model_out, deterministic=True)
deterministic_actions = squashed_deterministic_actions if config[
"normalize_actions"] else unsquash_actions(
squashed_deterministic_actions, action_space)
policy.output_actions, policy.sampled_action_logp = \
policy.exploration.get_exploration_action(
distribution_inputs, action_dist_class, model, explore, timestep)
actions = tf.cond(
tf.constant(explore) if isinstance(explore, bool) else explore,
true_fn=lambda: stochastic_actions,
false_fn=lambda: deterministic_actions)
logp = tf.cond(
tf.constant(explore) if isinstance(explore, bool) else explore,
true_fn=lambda: log_pis,
false_fn=lambda: tf.zeros_like(log_pis))
policy.output_actions, policy.action_logp = actions, logp
return policy.output_actions, policy.action_logp
return policy.output_actions, policy.sampled_action_logp
def actor_critic_loss(policy, model, _, train_batch):
@@ -156,9 +136,17 @@ def actor_critic_loss(policy, model, _, train_batch):
"obs": train_batch[SampleBatch.NEXT_OBS],
"is_training": policy._get_is_training_placeholder(),
}, [], None)
# TODO(hartikainen): figure actions and log pis
policy_t, log_pis_t = model.get_policy_output(model_out_t)
policy_tp1, log_pis_tp1 = model.get_policy_output(model_out_tp1)
action_dist_class = get_dist_class(policy.config, policy.action_space)
action_dist_t = action_dist_class(
model.action_model(model_out_t), policy.model)
policy_t = action_dist_t.sample()
log_pis_t = tf.expand_dims(action_dist_t.sampled_action_logp(), -1)
action_dist_tp1 = action_dist_class(
model.action_model(model_out_tp1), policy.model)
policy_tp1 = action_dist_tp1.sample()
log_pis_tp1 = tf.expand_dims(action_dist_tp1.sampled_action_logp(), -1)
log_alpha = model.log_alpha
alpha = model.alpha
rllib/tuned_examples/halfcheetah-sac.yaml
-1
View File
@@ -23,7 +23,6 @@ halfcheetah_sac:
target_network_update_freq: 1
timesteps_per_iteration: 1000
learning_starts: 10000
explore: True
optimization:
actor_learning_rate: 0.0003
critic_learning_rate: 0.0003
rllib/tuned_examples/pendulum-sac.yaml
-1
View File
@@ -24,7 +24,6 @@ pendulum_sac:
target_network_update_freq: 1
timesteps_per_iteration: 1000
learning_starts: 256
explore: True
optimization:
actor_learning_rate: 0.0003
critic_learning_rate: 0.0003