diff --git a/rllib/BUILD b/rllib/BUILD index 99ab84d21..f827ea40c 100644 --- a/rllib/BUILD +++ b/rllib/BUILD @@ -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"] ) diff --git a/rllib/agents/sac/__init__.py b/rllib/agents/sac/__init__.py index a9ea95da6..673192004 100644 --- a/rllib/agents/sac/__init__.py +++ b/rllib/agents/sac/__init__.py @@ -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", ] diff --git a/rllib/agents/sac/sac.py b/rllib/agents/sac/sac.py index 839cd1de1..91dfbfe4d 100644 --- a/rllib/agents/sac/sac.py +++ b/rllib/agents/sac/sac.py @@ -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 === diff --git a/rllib/agents/sac/sac_model.py b/rllib/agents/sac/sac_model.py index a7919f1f8..3c0ff5647 100644 --- a/rllib/agents/sac/sac_model.py +++ b/rllib/agents/sac/sac_model.py @@ -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.""" diff --git a/rllib/agents/sac/sac_policy.py b/rllib/agents/sac/sac_policy.py index 1b3eccb2b..e87124e14 100644 --- a/rllib/agents/sac/sac_policy.py +++ b/rllib/agents/sac/sac_policy.py @@ -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 diff --git a/rllib/tuned_examples/halfcheetah-sac.yaml b/rllib/tuned_examples/halfcheetah-sac.yaml index 81aaddda9..9fd836b29 100644 --- a/rllib/tuned_examples/halfcheetah-sac.yaml +++ b/rllib/tuned_examples/halfcheetah-sac.yaml @@ -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 diff --git a/rllib/tuned_examples/pendulum-sac.yaml b/rllib/tuned_examples/pendulum-sac.yaml index 9b320fb78..a2af901d0 100644 --- a/rllib/tuned_examples/pendulum-sac.yaml +++ b/rllib/tuned_examples/pendulum-sac.yaml @@ -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