diff --git a/ci/travis/install-dependencies.sh b/ci/travis/install-dependencies.sh index 8d319e1b6..542f183c5 100755 --- a/ci/travis/install-dependencies.sh +++ b/ci/travis/install-dependencies.sh @@ -91,11 +91,14 @@ else exit 1 fi +# Install modules needed in all jobs. +pip install dm-tree + # Additional RLlib dependencies. if [[ "$RLLIB_TESTING" == "1" ]]; then pip install tensorflow-probability==$tfp_version gast==0.2.2 \ torch==$torch_version torchvision \ - gym[atari] atari_py smart_open lz4 + atari_py gym[atari] lz4 smart_open fi # Additional streaming dependencies. diff --git a/doc/source/rllib-env.rst b/doc/source/rllib-env.rst index b2d959e7d..13ae346e7 100644 --- a/doc/source/rllib-env.rst +++ b/doc/source/rllib-env.rst @@ -19,7 +19,7 @@ DQN, Rainbow **Yes** `+parametric`_ No **Yes** DDPG, TD3 No **Yes** **Yes** APEX-DQN **Yes** `+parametric`_ No **Yes** APEX-DDPG No **Yes** **Yes** -SAC (todo) **Yes** **Yes** +SAC **Yes** **Yes** **Yes** ES **Yes** **Yes** No ARS **Yes** **Yes** No QMIX **Yes** No **Yes** `+RNN`_ diff --git a/python/setup.py b/python/setup.py index 5a014fa88..8e8ef7509 100644 --- a/python/setup.py +++ b/python/setup.py @@ -79,10 +79,12 @@ extras = { } extras["rllib"] = extras["tune"] + [ - "pyyaml", + "atari_py", + "dm_tree", "gym[atari]", - "opencv-python-headless", "lz4", + "opencv-python-headless", + "pyyaml", "scipy", ] diff --git a/rllib/BUILD b/rllib/BUILD index 96e3ec611..110bbbbf6 100644 --- a/rllib/BUILD +++ b/rllib/BUILD @@ -107,6 +107,14 @@ py_test( "agents/ppo/tests/test.py"] # TODO(sven): Move down once PR 6889 merged ) +# SAC +py_test( + name = "test_sac", + tags = ["agents_dir"], + size = "medium", + srcs = ["agents/sac/tests/test_sac.py"] +) + # TD3Trainer py_test( name = "test_td3", diff --git a/rllib/agents/dqn/dqn.py b/rllib/agents/dqn/dqn.py index ba96b5d02..1b7330c44 100644 --- a/rllib/agents/dqn/dqn.py +++ b/rllib/agents/dqn/dqn.py @@ -268,7 +268,7 @@ def get_initial_state(config): } -# TODO(sven): Move this to generic Trainer/Policy. Every Algo should do this. +# TODO(sven): Move this to generic Trainer. Every Algo should do this. def update_worker_exploration(trainer): """Sets epsilon exploration values in all policies to updated values. diff --git a/rllib/agents/dqn/tests/test_dqn.py b/rllib/agents/dqn/tests/test_dqn.py index 15e156ed8..533a2dd78 100644 --- a/rllib/agents/dqn/tests/test_dqn.py +++ b/rllib/agents/dqn/tests/test_dqn.py @@ -37,10 +37,12 @@ class TestDQN(unittest.TestCase): obs = np.array(0) # Test against all frameworks. - for fw in ["tf", "eager", "torch"]: + for fw in ["eager", "tf", "torch"]: if fw == "torch": continue + print("framework={}".format(fw)) + config["eager"] = True if fw == "eager" else False config["use_pytorch"] = True if fw == "torch" else False @@ -61,7 +63,7 @@ class TestDQN(unittest.TestCase): # (but no epsilon exploration). config["exploration_config"] = { "type": "SoftQ", - "temperature": 0.0 + "temperature": 0.001 } trainer = dqn.DQNTrainer(config=config, env="FrozenLake-v0") # Due to the low temp, always expect the same action. diff --git a/rllib/agents/sac/README.md b/rllib/agents/sac/README.md index f0c8e3315..7b600ec1f 100644 --- a/rllib/agents/sac/README.md +++ b/rllib/agents/sac/README.md @@ -1 +1,8 @@ -Implementation of Soft Actor-Critic (https://arxiv.org/abs/1812.05905.pdf). +Implementation of the Soft Actor-Critic algorithm: + +[1] Soft Actor-Critic Algorithms and Applications - T. Haarnoja, A. Zhou, K. Hartikainen, et. al +https://arxiv.org/abs/1812.05905.pdf + +For supporting discrete action spaces, we implemented this patch on top of the original algorithm: +[2] Soft Actor-Critic for Discrete Action Settings - Petros Christodoulou +https://arxiv.org/pdf/1910.07207v2.pdf diff --git a/rllib/agents/sac/sac.py b/rllib/agents/sac/sac.py index ed7f8d89b..d4f54c5d8 100644 --- a/rllib/agents/sac/sac.py +++ b/rllib/agents/sac/sac.py @@ -14,7 +14,6 @@ DEFAULT_CONFIG = with_common_config({ # === Model === "twin_q": True, "use_state_preprocessor": False, - "policy": "GaussianLatentSpacePolicy", # RLlib model options for the Q function "Q_model": { "hidden_activation": "relu", @@ -26,13 +25,17 @@ DEFAULT_CONFIG = with_common_config({ "hidden_layer_sizes": (256, 256), }, # Unsquash actions to the upper and lower bounds of env's action space. + # Ignored for discrete action spaces. "normalize_actions": True, # === Learning === # Update the target by \tau * policy + (1-\tau) * target_policy. "tau": 5e-3, - # Target entropy lower bound. This is the inverse of reward scale, - # and will be optimized automatically. + # Initial value to use for the entropy weight alpha. + "initial_alpha": 1.0, + # Target entropy lower bound. If "auto", will be set to -|A| (e.g. -2.0 for + # Discrete(2), -3.0 for Box(shape=(3,))). + # This is the inverse of reward scale, and will be optimized automatically. "target_entropy": "auto", # Disable setting done=True at end of episode. "no_done_at_end": True, @@ -47,13 +50,13 @@ DEFAULT_CONFIG = with_common_config({ # each worker will have a replay buffer of this size. "buffer_size": int(1e6), # If True prioritized replay buffer will be used. - # TODO(hartikainen): Make sure this works or remove the option. "prioritized_replay": False, "prioritized_replay_alpha": 0.6, "prioritized_replay_beta": 0.4, "prioritized_replay_eps": 1e-6, "prioritized_replay_beta_annealing_timesteps": 20000, "final_prioritized_replay_beta": 0.4, + "compress_observations": False, # === Optimization === @@ -62,7 +65,7 @@ DEFAULT_CONFIG = with_common_config({ "critic_learning_rate": 3e-4, "entropy_learning_rate": 3e-4, }, - # If not None, clip gradients during optimization at this value + # If not None, clip gradients during optimization at this value. "grad_norm_clipping": None, # How many steps of the model to sample before learning starts. "learning_starts": 1500, @@ -96,4 +99,7 @@ DEFAULT_CONFIG = with_common_config({ # yapf: enable SACTrainer = GenericOffPolicyTrainer.with_updates( - name="SAC", default_config=DEFAULT_CONFIG, default_policy=SACTFPolicy) + name="SAC", + default_config=DEFAULT_CONFIG, + default_policy=SACTFPolicy, +) diff --git a/rllib/agents/sac/sac_model.py b/rllib/agents/sac/sac_model.py index 3c0ff5647..f67341fcf 100644 --- a/rllib/agents/sac/sac_model.py +++ b/rllib/agents/sac/sac_model.py @@ -1,35 +1,14 @@ +from gym.spaces import Discrete import numpy as np from ray.rllib.models.tf.tf_modelv2 import TFModelV2 -from ray.rllib.utils import try_import_tf, try_import_tfp +from ray.rllib.utils import try_import_tf tf = try_import_tf() -tfp = try_import_tfp() SCALE_DIAG_MIN_MAX = (-20, 2) -def SquashBijector(): - # lazy def since it depends on tfp - class SquashBijector(tfp.bijectors.Bijector): - def __init__(self, validate_args=False, name="tanh"): - super(SquashBijector, self).__init__( - forward_min_event_ndims=0, - validate_args=validate_args, - name=name) - - def _forward(self, x): - return tf.nn.tanh(x) - - def _inverse(self, y): - return tf.atanh(y) - - def _forward_log_det_jacobian(self, x): - return 2. * (np.log(2.) - x - tf.nn.softplus(-2. * x)) - - return SquashBijector() - - class SACModel(TFModelV2): """Extension of standard TFModel for SAC. @@ -52,7 +31,8 @@ class SACModel(TFModelV2): actor_hiddens=(256, 256), critic_hidden_activation="relu", critic_hiddens=(256, 256), - twin_q=False): + twin_q=False, + initial_alpha=1.0): """Initialize variables of this model. Extra model kwargs: @@ -60,19 +40,26 @@ class SACModel(TFModelV2): actor_hiddens (list): hidden layers sizes for actor network critic_hidden_activation (str): activation for critic network critic_hiddens (list): hidden layers sizes for critic network - twin_q (bool): build twin Q networks + twin_q (bool): build twin Q networks. + initial_alpha (float): The initial value for the to-be-optimized + alpha parameter (default: 1.0). Note that the core layers for forward() are not defined here, this only defines the layers for the output heads. Those layers for forward() should be defined in subclasses of SACModel. """ - - if tfp is None: - raise ImportError("tensorflow-probability package not found") - super(SACModel, self).__init__(obs_space, action_space, num_outputs, model_config, name) - self.action_dim = np.product(action_space.shape) + self.discrete = False + if isinstance(action_space, Discrete): + self.action_dim = action_space.n + self.discrete = True + action_outs = q_outs = self.action_dim + else: + self.action_dim = np.product(action_space.shape) + action_outs = 2 * self.action_dim + q_outs = 1 + self.model_out = tf.keras.layers.Input( shape=(num_outputs, ), name="model_out") self.action_model = tf.keras.Sequential([ @@ -83,32 +70,39 @@ class SACModel(TFModelV2): for i, hidden in enumerate(actor_hiddens) ] + [ tf.keras.layers.Dense( - units=2 * self.action_dim, activation=None, name="action_out") + units=action_outs, activation=None, name="action_out") ]) self.shift_and_log_scale_diag = self.action_model(self.model_out) self.register_variables(self.action_model.variables) - self.actions_input = tf.keras.layers.Input( - shape=(self.action_dim, ), name="actions") + self.actions_input = None + if not self.discrete: + self.actions_input = tf.keras.layers.Input( + shape=(self.action_dim, ), name="actions") def build_q_net(name, observations, actions): - q_net = tf.keras.Sequential([ + # For continuous actions: Feed obs and actions (concatenated) + # through the NN. For discrete actions, only obs. + q_net = tf.keras.Sequential(([ tf.keras.layers.Concatenate(axis=1), - ] + [ + ] if not self.discrete else []) + [ tf.keras.layers.Dense( units=units, - activation=getattr(tf.nn, critic_hidden_activation), + activation=getattr(tf.nn, critic_hidden_activation, None), name="{}_hidden_{}".format(name, i)) for i, units in enumerate(critic_hiddens) ] + [ tf.keras.layers.Dense( - units=1, activation=None, name="{}_out".format(name)) + units=q_outs, activation=None, name="{}_out".format(name)) ]) - # TODO(hartikainen): Remove the unnecessary Model call here - q_net = tf.keras.Model([observations, actions], - q_net([observations, actions])) + # TODO(hartikainen): Remove the unnecessary Model calls here + if self.discrete: + q_net = tf.keras.Model(observations, q_net(observations)) + else: + q_net = tf.keras.Model([observations, actions], + q_net([observations, actions])) return q_net self.q_net = build_q_net("q", self.model_out, self.actions_input) @@ -121,12 +115,13 @@ class SACModel(TFModelV2): else: self.twin_q_net = None - self.log_alpha = tf.Variable(0.0, dtype=tf.float32, name="log_alpha") + self.log_alpha = tf.Variable( + np.log(initial_alpha), dtype=tf.float32, name="log_alpha") self.alpha = tf.exp(self.log_alpha) self.register_variables([self.log_alpha]) - def get_q_values(self, model_out, actions): + def get_q_values(self, model_out, actions=None): """Return the Q estimates for the most recent forward pass. This implements Q(s, a). @@ -134,16 +129,19 @@ class SACModel(TFModelV2): Arguments: model_out (Tensor): obs embeddings from the model layers, of shape [BATCH_SIZE, num_outputs]. - actions (Tensor): action values that correspond with the most - recent batch of observations passed through forward(), of shape - [BATCH_SIZE, action_dim]. + actions (Optional[Tensor]): Actions to return the Q-values for. + Shape: [BATCH_SIZE, action_dim]. If None (discrete action + case), return Q-values for all actions. Returns: tensor of shape [BATCH_SIZE]. """ - return self.q_net([model_out, actions]) + if actions is not None: + return self.q_net([model_out, actions]) + else: + return self.q_net(model_out) - def get_twin_q_values(self, model_out, actions): + def get_twin_q_values(self, model_out, actions=None): """Same as get_q_values but using the twin Q net. This implements the twin Q(s, a). @@ -151,14 +149,17 @@ class SACModel(TFModelV2): Arguments: model_out (Tensor): obs embeddings from the model layers, of shape [BATCH_SIZE, num_outputs]. - actions (Tensor): action values that correspond with the most - recent batch of observations passed through forward(), of shape - [BATCH_SIZE, action_dim]. + actions (Optional[Tensor]): Actions to return the Q-values for. + Shape: [BATCH_SIZE, action_dim]. If None (discrete action + case), return Q-values for all actions. Returns: tensor of shape [BATCH_SIZE]. """ - return self.twin_q_net([model_out, actions]) + if actions is not None: + return self.twin_q_net([model_out, actions]) + else: + return self.twin_q_net(model_out) def policy_variables(self): """Return the list of variables for the policy net.""" diff --git a/rllib/agents/sac/sac_policy.py b/rllib/agents/sac/sac_policy.py index ff8883d5b..468c5ae27 100644 --- a/rllib/agents/sac/sac_policy.py +++ b/rllib/agents/sac/sac_policy.py @@ -1,6 +1,6 @@ -from gym.spaces import Box -import numpy as np +from gym.spaces import Box, Discrete import logging +import numpy as np import ray import ray.experimental.tf_utils @@ -12,10 +12,11 @@ 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.models.tf.tf_action_dist import ( + Categorical, SquashedGaussian, DiagGaussian) from ray.rllib.utils import try_import_tf, try_import_tfp from ray.rllib.utils.annotations import override +from ray.rllib.utils.error import UnsupportedSpaceException from ray.rllib.utils.tf_ops import minimize_and_clip, make_tf_callable tf = try_import_tf() @@ -30,10 +31,10 @@ def build_sac_model(policy, obs_space, action_space, config): "Setting use_state_preprocessor=True since a custom model " "was specified.") config["use_state_preprocessor"] = True - if not isinstance(action_space, Box): + if not isinstance(action_space, (Box, Discrete)): raise UnsupportedSpaceException( "Action space {} is not supported for SAC.".format(action_space)) - if len(action_space.shape) > 1: + if isinstance(action_space, Box) and len(action_space.shape) > 1: raise UnsupportedSpaceException( "Action space has multiple dimensions " "{}. ".format(action_space.shape) + @@ -61,7 +62,8 @@ def build_sac_model(policy, obs_space, action_space, config): actor_hiddens=config["policy_model"]["hidden_layer_sizes"], critic_hidden_activation=config["Q_model"]["hidden_activation"], critic_hiddens=config["Q_model"]["hidden_layer_sizes"], - twin_q=config["twin_q"]) + twin_q=config["twin_q"], + initial_alpha=config["initial_alpha"]) policy.target_model = ModelCatalog.get_model_v2( obs_space, @@ -76,7 +78,8 @@ def build_sac_model(policy, obs_space, action_space, config): actor_hiddens=config["policy_model"]["hidden_layer_sizes"], critic_hidden_activation=config["Q_model"]["hidden_activation"], critic_hiddens=config["Q_model"]["hidden_layer_sizes"], - twin_q=config["twin_q"]) + twin_q=config["twin_q"], + initial_alpha=config["initial_alpha"]) return policy.model @@ -89,8 +92,12 @@ def postprocess_trajectory(policy, def get_dist_class(config, action_space): - action_dist_class = SquashedGaussian if \ - config["normalize_actions"] is True else DiagGaussian + if isinstance(action_space, Discrete): + action_dist_class = Categorical + else: + action_dist_class = ( + SquashedGaussian if config["normalize_actions"] + else DiagGaussian) return action_dist_class @@ -137,48 +144,78 @@ def actor_critic_loss(policy, model, _, train_batch): "is_training": policy._get_is_training_placeholder(), }, [], None) - 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) + # Discrete case. + if model.discrete: + # Get all action probs directly from pi and form their logp. + log_pis_t = tf.nn.log_softmax(model.action_model(model_out_t), -1) + policy_t = tf.exp(log_pis_t) + log_pis_tp1 = tf.nn.log_softmax(model.action_model(model_out_tp1), -1) + policy_tp1 = tf.exp(log_pis_tp1) + # Q-values. + q_t = model.get_q_values(model_out_t) + # Target Q-values. + q_tp1 = policy.target_model.get_q_values(target_model_out_tp1) + if policy.config["twin_q"]: + twin_q_t = model.get_twin_q_values(model_out_t) + twin_q_tp1 = policy.target_model.get_twin_q_values( + target_model_out_tp1) + q_tp1 = tf.reduce_min((q_tp1, twin_q_tp1), axis=0) + q_tp1 -= model.alpha * log_pis_tp1 - 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) + # Actually selected Q-values (from the actions batch). + one_hot = tf.one_hot( + train_batch[SampleBatch.ACTIONS], depth=q_t.shape.as_list()[-1]) + q_t_selected = tf.reduce_sum(q_t * one_hot, axis=-1) + if policy.config["twin_q"]: + twin_q_t_selected = tf.reduce_sum(twin_q_t * one_hot, axis=-1) + # Discrete case: "Best" means weighted by the policy (prob) outputs. + q_tp1_best = tf.reduce_sum(tf.multiply(policy_tp1, q_tp1), axis=-1) + q_tp1_best_masked = \ + (1.0 - tf.cast(train_batch[SampleBatch.DONES], tf.float32)) * \ + q_tp1_best + # Continuous actions case. + else: + # Sample simgle actions from distribution. + 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 + # Q-values for the actually selected actions. + q_t = model.get_q_values(model_out_t, train_batch[SampleBatch.ACTIONS]) + if policy.config["twin_q"]: + twin_q_t = model.get_twin_q_values( + model_out_t, train_batch[SampleBatch.ACTIONS]) - # q network evaluation - q_t = model.get_q_values(model_out_t, train_batch[SampleBatch.ACTIONS]) - if policy.config["twin_q"]: - twin_q_t = model.get_twin_q_values(model_out_t, - train_batch[SampleBatch.ACTIONS]) + # Q-values for current policy in given current state. + q_t_det_policy = model.get_q_values(model_out_t, policy_t) + if policy.config["twin_q"]: + twin_q_t_det_policy = model.get_twin_q_values( + model_out_t, policy_t) + q_t_det_policy = tf.reduce_min( + (q_t_det_policy, twin_q_t_det_policy), axis=0) - # Q-values for current policy (no noise) in given current state - q_t_det_policy = model.get_q_values(model_out_t, policy_t) - if policy.config["twin_q"]: - twin_q_t_det_policy = model.get_twin_q_values(model_out_t, policy_t) - q_t_det_policy = tf.reduce_min( - (q_t_det_policy, twin_q_t_det_policy), axis=0) + # target q network evaluation + q_tp1 = policy.target_model.get_q_values(target_model_out_tp1, + policy_tp1) + if policy.config["twin_q"]: + twin_q_tp1 = policy.target_model.get_twin_q_values( + target_model_out_tp1, policy_tp1) - # target q network evaluation - q_tp1 = policy.target_model.get_q_values(target_model_out_tp1, policy_tp1) - if policy.config["twin_q"]: - twin_q_tp1 = policy.target_model.get_twin_q_values( - target_model_out_tp1, policy_tp1) + q_t_selected = tf.squeeze(q_t, axis=len(q_t.shape) - 1) + if policy.config["twin_q"]: + twin_q_t_selected = tf.squeeze(twin_q_t, axis=len(q_t.shape) - 1) + q_tp1 = tf.reduce_min((q_tp1, twin_q_tp1), axis=0) + q_tp1 -= model.alpha * log_pis_tp1 - q_t_selected = tf.squeeze(q_t, axis=len(q_t.shape) - 1) - if policy.config["twin_q"]: - twin_q_t_selected = tf.squeeze(twin_q_t, axis=len(q_t.shape) - 1) - q_tp1 = tf.reduce_min((q_tp1, twin_q_tp1), axis=0) - q_tp1 -= alpha * log_pis_tp1 - - q_tp1_best = tf.squeeze(input=q_tp1, axis=len(q_tp1.shape) - 1) - q_tp1_best_masked = ( - 1.0 - tf.cast(train_batch[SampleBatch.DONES], tf.float32)) * q_tp1_best + q_tp1_best = tf.squeeze(input=q_tp1, axis=len(q_tp1.shape) - 1) + q_tp1_best_masked = (1.0 - tf.cast(train_batch[SampleBatch.DONES], + tf.float32)) * q_tp1_best assert policy.config["n_step"] == 1, "TODO(hartikainen) n_step > 1" @@ -187,13 +224,13 @@ def actor_critic_loss(policy, model, _, train_batch): train_batch[SampleBatch.REWARDS] + policy.config["gamma"]**policy.config["n_step"] * q_tp1_best_masked) - # compute the error (potentially clipped) + # Compute the TD-error (potentially clipped). + base_td_error = tf.abs(q_t_selected - q_t_selected_target) if policy.config["twin_q"]: - base_td_error = q_t_selected - q_t_selected_target - twin_td_error = twin_q_t_selected - q_t_selected_target - td_error = 0.5 * (tf.square(base_td_error) + tf.square(twin_td_error)) + twin_td_error = tf.abs(twin_q_t_selected - q_t_selected_target) + td_error = 0.5 * (base_td_error + twin_td_error) else: - td_error = tf.square(q_t_selected - q_t_selected_target) + td_error = base_td_error critic_loss = [ tf.losses.mean_squared_error( @@ -206,12 +243,38 @@ def actor_critic_loss(policy, model, _, train_batch): predictions=twin_q_t_selected, weights=0.5)) - target_entropy = (-np.prod(policy.action_space.shape) - if policy.config["target_entropy"] == "auto" else - policy.config["target_entropy"]) - alpha_loss = -tf.reduce_mean( - log_alpha * tf.stop_gradient(log_pis_t + target_entropy)) - actor_loss = tf.reduce_mean(alpha * log_pis_t - q_t_det_policy) + # Auto-calculate the target entropy. + if policy.config["target_entropy"] == "auto": + if model.discrete: + target_entropy = np.array(-policy.action_space.n, dtype=np.float32) + else: + target_entropy = -np.prod(policy.action_space.shape) + else: + target_entropy = policy.config["target_entropy"] + + # Alpha- and actor losses. + # Note: In the papers, alpha is used directly, here we take the log. + # Discrete case: Multiply the action probs as weights with the original + # loss terms (no expectations needed). + if model.discrete: + alpha_loss = tf.reduce_mean( + tf.reduce_sum( + tf.multiply( + tf.stop_gradient(policy_t), -model.log_alpha * + tf.stop_gradient(log_pis_t + target_entropy)), + axis=-1)) + actor_loss = tf.reduce_mean( + tf.reduce_sum( + tf.multiply( + # NOTE: No stop_grad around policy output here + # (compare with q_t_det_policy for continuous case). + policy_t, + model.alpha * log_pis_t - tf.stop_gradient(q_t)), + axis=-1)) + else: + alpha_loss = -tf.reduce_mean( + model.log_alpha * tf.stop_gradient(log_pis_t + target_entropy)) + actor_loss = tf.reduce_mean(model.alpha * log_pis_t - q_t_det_policy) # save for stats function policy.q_t = q_t @@ -219,6 +282,8 @@ def actor_critic_loss(policy, model, _, train_batch): policy.actor_loss = actor_loss policy.critic_loss = critic_loss policy.alpha_loss = alpha_loss + policy.alpha_value = model.alpha + policy.target_entropy = target_entropy # in a custom apply op we handle the losses separately, but return them # combined in one loss for now @@ -226,7 +291,7 @@ def actor_critic_loss(policy, model, _, train_batch): def gradients(policy, optimizer, loss): - if policy.config["grad_norm_clipping"] is not None: + if policy.config["grad_norm_clipping"]: actor_grads_and_vars = minimize_and_clip( optimizer, policy.actor_loss, @@ -315,6 +380,9 @@ def stats(policy, train_batch): "td_error": tf.reduce_mean(policy.td_error), "actor_loss": tf.reduce_mean(policy.actor_loss), "critic_loss": tf.reduce_mean(policy.critic_loss), + "alpha_loss": tf.reduce_mean(policy.alpha_loss), + "alpha_value": tf.reduce_mean(policy.alpha_value), + "target_entropy": tf.constant(policy.target_entropy), "mean_q": tf.reduce_mean(policy.q_t), "max_q": tf.reduce_max(policy.q_t), "min_q": tf.reduce_min(policy.q_t), @@ -346,7 +414,8 @@ class ComputeTDErrorMixin: @make_tf_callable(self.get_session(), dynamic_shape=True) def compute_td_error(obs_t, act_t, rew_t, obs_tp1, done_mask, importance_weights): - # Do forward pass on loss to update td error attribute + # Do forward pass on loss to update td errors attribute + # (one TD-error value per item in batch to update PR weights). actor_critic_loss( self, self.model, None, { SampleBatch.CUR_OBS: tf.convert_to_tensor(obs_t), diff --git a/rllib/agents/sac/tests/test_sac.py b/rllib/agents/sac/tests/test_sac.py new file mode 100644 index 000000000..771b264b3 --- /dev/null +++ b/rllib/agents/sac/tests/test_sac.py @@ -0,0 +1,39 @@ +import unittest + +import ray +import ray.rllib.agents.sac as sac +from ray.rllib.utils.framework import try_import_tf + +tf = try_import_tf() + + +class TestSAC(unittest.TestCase): + def test_sac_compilation(self): + """Test whether an SACTrainer can be built with all frameworks.""" + ray.init() + config = sac.DEFAULT_CONFIG.copy() + config["num_workers"] = 0 # Run locally. + num_iterations = 1 + + # eager (discrete and cont. actions). + for fw in ["eager", "tf", "torch"]: + print("framework={}".format(fw)) + if fw == "torch": + continue + config["eager"] = fw == "eager" + config["use_pytorch"] = fw == "torch" + for env in [ + "CartPole-v0", + "Pendulum-v0", + ]: + print("Env={}".format(env)) + trainer = sac.SACTrainer(config=config, env=env) + for i in range(num_iterations): + results = trainer.train() + print(results) + + +if __name__ == "__main__": + import pytest + import sys + sys.exit(pytest.main(["-v", __file__])) diff --git a/rllib/models/tests/test_distributions.py b/rllib/models/tests/test_distributions.py index 3d7f3adf1..b3f419420 100644 --- a/rllib/models/tests/test_distributions.py +++ b/rllib/models/tests/test_distributions.py @@ -5,7 +5,7 @@ from tensorflow.python.eager.context import eager_mode import unittest from ray.rllib.models.tf.tf_action_dist import Categorical, MultiCategorical, \ - SquashedGaussian + SquashedGaussian, GumbelSoftmax from ray.rllib.models.torch.torch_action_dist import TorchMultiCategorical from ray.rllib.utils import try_import_tf, try_import_torch from ray.rllib.utils.numpy import MIN_LOG_NN_OUTPUT, MAX_LOG_NN_OUTPUT, softmax @@ -50,7 +50,6 @@ class TestDistributions(unittest.TestCase): shape=(num_sub_distributions, batch_size), dtype=np.int32) - # The Component to test. inputs = inputs_space.sample() input_lengths = [num_categories] * num_sub_distributions inputs_split = np.split(inputs, num_sub_distributions, axis=1) @@ -58,6 +57,7 @@ class TestDistributions(unittest.TestCase): for fw in ["tf", "eager", "torch"]: print("framework={}".format(fw)) + # Create the correct distribution object. cls = MultiCategorical if fw != "torch" else TorchMultiCategorical multi_categorical = cls(inputs, None, input_lengths) @@ -169,6 +169,30 @@ class TestDistributions(unittest.TestCase): out = squashed_distribution.logp(values) check(out, log_prob) + def test_gumbel_softmax(self): + """Tests the GumbelSoftmax ActionDistribution (tf-eager only).""" + with eager_mode(): + batch_size = 1000 + num_categories = 5 + input_space = Box(-1.0, 1.0, shape=(batch_size, num_categories)) + + # Batch of size=n and deterministic. + inputs = input_space.sample() + gumbel_softmax = GumbelSoftmax(inputs, {}, temperature=1.0) + + expected = softmax(inputs) + # Sample n times, expect always mean value (deterministic draw). + out = gumbel_softmax.deterministic_sample() + check(out, expected) + + # Batch of size=n and non-deterministic -> expect roughly that + # the max-likelihood (argmax) ints are output (most of the time). + inputs = input_space.sample() + gumbel_softmax = GumbelSoftmax(inputs, {}, temperature=1.0) + expected_mean = np.mean(np.argmax(inputs, -1)).astype(np.float32) + outs = gumbel_softmax.sample() + check(np.mean(np.argmax(outs, -1)), expected_mean, rtol=0.08) + if __name__ == "__main__": import unittest diff --git a/rllib/models/tf/tf_action_dist.py b/rllib/models/tf/tf_action_dist.py index c6080fc52..3042e16b1 100644 --- a/rllib/models/tf/tf_action_dist.py +++ b/rllib/models/tf/tf_action_dist.py @@ -45,7 +45,8 @@ class Categorical(TFActionDistribution): @DeveloperAPI def __init__(self, inputs, model=None, temperature=1.0): - temperature = max(0.0001, temperature) # clamp for stability reasons + assert temperature > 0.0, "Categorical `temperature` must be > 0.0!" + self.n = inputs.shape[-1] # Allow softmax formula w/ temperature != 1.0: # Divide inputs by temperature. super().__init__(inputs / temperature, model) @@ -143,6 +144,68 @@ class MultiCategorical(TFActionDistribution): return np.sum(action_space.nvec) +class GumbelSoftmax(TFActionDistribution): + """GumbelSoftmax distr. (for differentiable sampling in discr. actions + + The Gumbel Softmax distribution [1] (also known as the Concrete [2] + distribution) is a close cousin of the relaxed one-hot categorical + distribution, whose tfp implementation we will use here plus + adjusted `sample_...` and `log_prob` methods. See discussion at [0]. + + [0] https://stackoverflow.com/questions/56226133/ + soft-actor-critic-with-discrete-action-space + + [1] Categorical Reparametrization with Gumbel-Softmax (Jang et al, 2017): + https://arxiv.org/abs/1611.01144 + [2] The Concrete Distribution: A Continuous Relaxation of Discrete Random + Variables (Maddison et al, 2017) https://arxiv.org/abs/1611.00712 + """ + + @DeveloperAPI + def __init__(self, inputs, model=None, temperature=1.0): + """Initializes a GumbelSoftmax distribution. + + Args: + temperature (float): Temperature parameter. For low temperatures, + the expected value approaches a categorical random variable. + For high temperatures, the expected value approaches a uniform + distribution. + """ + assert temperature >= 0.0 + self.dist = tfp.distributions.RelaxedOneHotCategorical( + temperature=temperature, logits=inputs) + super().__init__(inputs, model) + + @override(ActionDistribution) + def deterministic_sample(self): + # Return the dist object's prob values. + return self.dist._distribution.probs + + @override(ActionDistribution) + def logp(self, x): + # Override since the implementation of tfp.RelaxedOneHotCategorical + # yields positive values. + if x.shape != self.dist.logits.shape: + values = tf.one_hot( + x, self.dist.logits.shape.as_list()[-1], dtype=tf.float32) + assert values.shape == self.dist.logits.shape, ( + values.shape, self.dist.logits.shape) + + # [0]'s implementation (see line below) seems to be an approximation + # to the actual Gumbel Softmax density. + return -tf.reduce_sum( + -x * tf.nn.log_softmax(self.dist.logits, axis=-1), axis=-1) + + @override(TFActionDistribution) + def _build_sample_op(self): + return self.dist.sample() + + @staticmethod + @override(ActionDistribution) + def required_model_output_shape(action_space, model_config): + return action_space.n + + class DiagGaussian(TFActionDistribution): """Action distribution where each vector element is a gaussian. diff --git a/rllib/models/torch/torch_action_dist.py b/rllib/models/torch/torch_action_dist.py index e63ea8beb..f3d6f03e7 100644 --- a/rllib/models/torch/torch_action_dist.py +++ b/rllib/models/torch/torch_action_dist.py @@ -44,9 +44,11 @@ class TorchCategorical(TorchDistributionWrapper): """Wrapper class for PyTorch Categorical distribution.""" @override(ActionDistribution) - def __init__(self, inputs, model): - super().__init__(inputs, model) - self.dist = torch.distributions.categorical.Categorical(logits=inputs) + def __init__(self, inputs, model=None, temperature=1.0): + assert temperature > 0.0, "Categorical `temperature` must be > 0.0!" + super().__init__(inputs / temperature, model) + self.dist = torch.distributions.categorical.Categorical( + logits=self.inputs) @override(ActionDistribution) def deterministic_sample(self): diff --git a/rllib/policy/eager_tf_policy.py b/rllib/policy/eager_tf_policy.py index 3db374a1a..18eda5227 100644 --- a/rllib/policy/eager_tf_policy.py +++ b/rllib/policy/eager_tf_policy.py @@ -5,6 +5,7 @@ It supports both traced and non-traced eager execution modes.""" import logging import functools import numpy as np +import tree from ray.util.debug import log_once from ray.rllib.evaluation.episode import _flatten_action @@ -23,12 +24,12 @@ logger = logging.getLogger(__name__) def _convert_to_tf(x): if isinstance(x, SampleBatch): x = {k: v for k, v in x.items() if k != SampleBatch.INFOS} - return tf.nest.map_structure(_convert_to_tf, x) + return tree.map_structure(_convert_to_tf, x) if isinstance(x, Policy): return x if x is not None: - x = tf.nest.map_structure( + x = tree.map_structure( lambda f: tf.convert_to_tensor(f) if f is not None else None, x) return x @@ -37,7 +38,7 @@ def _convert_to_numpy(x): if x is None: return None try: - return tf.nest.map_structure(lambda component: component.numpy(), x) + return tree.map_structure(lambda component: component.numpy(), x) except AttributeError: raise TypeError( ("Object of type {} has no method to convert to numpy.").format( @@ -65,7 +66,7 @@ def convert_eager_outputs(func): def _func(*args, **kwargs): out = func(*args, **kwargs) if tf.executing_eagerly(): - out = tf.nest.map_structure(_convert_to_numpy, out) + out = tree.map_structure(_convert_to_numpy, out) return out return _func @@ -545,18 +546,14 @@ def build_eager_tf_policy(name, action_dtype, action_shape = ModelCatalog.get_action_shape( self.action_space) dummy_batch = { - SampleBatch.CUR_OBS: tf.convert_to_tensor( - np.array([self.observation_space.sample()])), - SampleBatch.NEXT_OBS: tf.convert_to_tensor( - np.array([self.observation_space.sample()])), - SampleBatch.DONES: tf.convert_to_tensor( - np.array([False], dtype=np.bool)), - SampleBatch.ACTIONS: tf.convert_to_tensor( - np.zeros( - (1, ) + action_shape[1:], - dtype=action_dtype.as_numpy_dtype())), - SampleBatch.REWARDS: tf.convert_to_tensor( - np.array([0], dtype=np.float32)), + SampleBatch.CUR_OBS: np.array( + [self.observation_space.sample()]), + SampleBatch.NEXT_OBS: np.array( + [self.observation_space.sample()]), + SampleBatch.DONES: np.array([False], dtype=np.bool), + SampleBatch.ACTIONS: tree.map_structure( + lambda c: np.array([c]), self.action_space.sample()), + SampleBatch.REWARDS: np.array([0], dtype=np.float32), } if obs_include_prev_action_reward: dummy_batch.update({ @@ -568,8 +565,10 @@ def build_eager_tf_policy(name, dummy_batch["state_out_{}".format(i)] = h if self._state_in: - dummy_batch["seq_lens"] = tf.convert_to_tensor( - np.array([1], dtype=np.int32)) + dummy_batch["seq_lens"] = np.array([1], dtype=np.int32) + + # Convert everything to tensors. + dummy_batch = tree.map_structure(tf.convert_to_tensor, dummy_batch) # for IMPALA which expects a certain sample batch size. def tile_to(tensor, n): @@ -577,10 +576,9 @@ def build_eager_tf_policy(name, [n] + [1 for _ in tensor.shape.as_list()[1:]]) if get_batch_divisibility_req: - dummy_batch = { - k: tile_to(v, get_batch_divisibility_req(self)) - for k, v in dummy_batch.items() - } + dummy_batch = tree.map_structure( + lambda c: tile_to(c, get_batch_divisibility_req(self)), + dummy_batch) # Execute a forward pass to get self.action_dist etc initialized, # and also obtain the extra action fetches @@ -597,10 +595,8 @@ def build_eager_tf_policy(name, # overwrite any tensor state from that call) self.model.from_batch(dummy_batch) - postprocessed_batch = { - k: tf.convert_to_tensor(v) - for k, v in postprocessed_batch.items() - } + postprocessed_batch = tree.map_structure( + lambda c: tf.convert_to_tensor(c), postprocessed_batch.data) loss_fn(self, self.model, self.dist_class, postprocessed_batch) if stats_fn: diff --git a/rllib/policy/policy.py b/rllib/policy/policy.py index 5093316d2..4e0b2d19d 100644 --- a/rllib/policy/policy.py +++ b/rllib/policy/policy.py @@ -120,7 +120,7 @@ class Policy(metaclass=ABCMeta): episode (MultiAgentEpisode): this provides access to all of the internal episode state, which may be useful for model-based or multi-agent algorithms. - clip_actions (bool): should the action be clipped + clip_actions (bool): Should actions be clipped? explore (bool): Whether to pick an exploitation or exploration action (default: None -> use self.config["explore"]). timestep (int): The current (sampling) time step. diff --git a/rllib/policy/tf_policy.py b/rllib/policy/tf_policy.py index d0fadd0bf..11f6b17f2 100644 --- a/rllib/policy/tf_policy.py +++ b/rllib/policy/tf_policy.py @@ -1,6 +1,7 @@ import errno import logging import os +import tree import numpy as np import ray @@ -492,7 +493,7 @@ class TFPolicy(Policy): # build output signatures output_signature = self._extra_output_signature_def() - for i, a in enumerate(tf.nest.flatten(self._sampled_action)): + for i, a in enumerate(tree.flatten(self._sampled_action)): output_signature["actions_{}".format(i)] = \ tf.saved_model.utils.build_tensor_info(a) diff --git a/rllib/tests/test_eager_support.py b/rllib/tests/test_eager_support.py index 6a02f6355..878df7499 100644 --- a/rllib/tests/test_eager_support.py +++ b/rllib/tests/test_eager_support.py @@ -15,11 +15,11 @@ def check_support(alg, config, test_trace=True): config["log_level"] = "ERROR" config["eager_tracing"] = False - tune.run(a, config=config, stop={"training_iteration": 0}) + tune.run(a, config=config, stop={"training_iteration": 1}) if test_trace: config["eager_tracing"] = True - tune.run(a, config=config, stop={"training_iteration": 0}) + tune.run(a, config=config, stop={"training_iteration": 1}) class TestEagerSupport(unittest.TestCase): diff --git a/rllib/tuned_examples/regression_tests/cartpole-sac.yaml b/rllib/tuned_examples/regression_tests/cartpole-sac.yaml new file mode 100644 index 000000000..830a923fb --- /dev/null +++ b/rllib/tuned_examples/regression_tests/cartpole-sac.yaml @@ -0,0 +1,19 @@ +cartpole-sac: + env: CartPole-v0 + run: SAC + stop: + episode_reward_mean: 150 + config: + gamma: 0.95 + no_done_at_end: false + target_network_update_freq: 32 + tau: 1.0 + # initial_alpha: 0.5 + train_batch_size: 32 + optimization: + actor_learning_rate: 0.005 + critic_learning_rate: 0.005 + entropy_learning_rate: 0.0001 + # grad_norm_clipping: 40.0 + evaluation_config: + explore: true diff --git a/rllib/utils/tf_run_builder.py b/rllib/utils/tf_run_builder.py index 5ee6b4fb9..5a90bcd7d 100644 --- a/rllib/utils/tf_run_builder.py +++ b/rllib/utils/tf_run_builder.py @@ -42,11 +42,10 @@ class TFRunBuilder: self._executed = run_timeline( self.session, self.fetches, self.debug_name, self.feed_dict, os.environ.get("TF_TIMELINE_DIR")) - except Exception: + except Exception as e: logger.exception("Error fetching: {}, feed_dict={}".format( self.fetches, self.feed_dict)) - raise ValueError("Error fetching: {}, feed_dict={}".format( - self.fetches, self.feed_dict)) + raise e if isinstance(to_fetch, int): return self._executed[to_fetch] elif isinstance(to_fetch, list):