From 510c8506512b8e9873ffbab0abdc1a668b022022 Mon Sep 17 00:00:00 2001 From: Sven Mika Date: Fri, 6 Mar 2020 19:37:12 +0100 Subject: [PATCH] [RLlib] SAC add discrete action support. (#7320) * Exploration API (+EpsilonGreedy sub-class). * Exploration API (+EpsilonGreedy sub-class). * Cleanup/LINT. * Add `deterministic` to generic Trainer config (NOTE: this is still ignored by most Agents). * Add `error` option to deprecation_warning(). * WIP. * Bug fix: Get exploration-info for tf framework. Bug fix: Properly deprecate some DQN config keys. * WIP. * LINT. * WIP. * Split PerWorkerEpsilonGreedy out of EpsilonGreedy. Docstrings. * Fix bug in sampler.py in case Policy has self.exploration = None * Update rllib/agents/dqn/dqn.py Co-Authored-By: Eric Liang * WIP. * Update rllib/agents/trainer.py Co-Authored-By: Eric Liang * WIP. * Change requests. * LINT * In tune/utils/util.py::deep_update() Only keep deep_updat'ing if both original and value are dicts. If value is not a dict, set * Completely obsolete syn_replay_optimizer.py's parameters schedule_max_timesteps AND beta_annealing_fraction (replaced with prioritized_replay_beta_annealing_timesteps). * Update rllib/evaluation/worker_set.py Co-Authored-By: Eric Liang * Review fixes. * Fix default value for DQN's exploration spec. * LINT * Fix recursion bug (wrong parent c'tor). * Do not pass timestep to get_exploration_info. * Update tf_policy.py * Fix some remaining issues with test cases and remove more deprecated DQN/APEX exploration configs. * Bug fix tf-action-dist * DDPG incompatibility bug fix with new DQN exploration handling (which is imported by DDPG). * Switch off exploration when getting action probs from off-policy-estimator's policy. * LINT * Fix test_checkpoint_restore.py. * Deprecate all SAC exploration (unused) configs. * Properly use `model.last_output()` everywhere. Instead of `model._last_output`. * WIP. * Take out set_epsilon from multi-agent-env test (not needed, decays anyway). * WIP. * Trigger re-test (flaky checkpoint-restore test). * WIP. * WIP. * Add test case for deterministic action sampling in PPO. * bug fix. * Added deterministic test cases for different Agents. * Fix problem with TupleActions in dynamic-tf-policy. * Separate supported_spaces tests so they can be run separately for easier debugging. * LINT. * Fix autoregressive_action_dist.py test case. * Re-test. * Fix. * Remove duplicate py_test rule from bazel. * LINT. * WIP. * WIP. * SAC fix. * SAC fix. * WIP. * WIP. * WIP. * FIX 2 examples tests. * WIP. * WIP. * WIP. * WIP. * WIP. * Fix. * LINT. * Renamed test file. * WIP. * Add unittest.main. * Make action_dist_class mandatory. * fix * FIX. * WIP. * WIP. * Fix. * Fix. * Fix explorations test case (contextlib cannot find its own nullcontext??). * Force torch to be installed for QMIX. * LINT. * Fix determine_tests_to_run.py. * Fix determine_tests_to_run.py. * WIP * Add Random exploration component to tests (fixed issue with "static-graph randomness" via py_function). * Add Random exploration component to tests (fixed issue with "static-graph randomness" via py_function). * Rename some stuff. * Rename some stuff. * WIP. * update. * WIP. * Gumbel Softmax Dist. * WIP. * WIP. * WIP. * WIP. * WIP. * WIP. * WIP * WIP. * WIP. * Hypertune. * Hypertune. * Hypertune. * Lock-in. * Cleanup. * LINT. * Fix. * Update rllib/policy/eager_tf_policy.py Co-Authored-By: Kristian Hartikainen * Update rllib/agents/sac/sac_policy.py Co-Authored-By: Kristian Hartikainen * Update rllib/agents/sac/sac_policy.py Co-Authored-By: Kristian Hartikainen * Update rllib/models/tf/tf_action_dist.py Co-Authored-By: Kristian Hartikainen * Update rllib/models/tf/tf_action_dist.py Co-Authored-By: Kristian Hartikainen * Fix items from review comments. * Add dm_tree to RLlib dependencies. * Add dm_tree to RLlib dependencies. * Fix DQN test cases ((Torch)Categorical). * Fix wrong pip install. Co-authored-by: Eric Liang Co-authored-by: Kristian Hartikainen --- ci/travis/install-dependencies.sh | 5 +- doc/source/rllib-env.rst | 2 +- python/setup.py | 6 +- rllib/BUILD | 8 + rllib/agents/dqn/dqn.py | 2 +- rllib/agents/dqn/tests/test_dqn.py | 6 +- rllib/agents/sac/README.md | 9 +- rllib/agents/sac/sac.py | 18 +- rllib/agents/sac/sac_model.py | 103 +++++----- rllib/agents/sac/sac_policy.py | 187 ++++++++++++------ rllib/agents/sac/tests/test_sac.py | 39 ++++ rllib/models/tests/test_distributions.py | 28 ++- rllib/models/tf/tf_action_dist.py | 65 +++++- rllib/models/torch/torch_action_dist.py | 8 +- rllib/policy/eager_tf_policy.py | 48 +++-- rllib/policy/policy.py | 2 +- rllib/policy/tf_policy.py | 3 +- rllib/tests/test_eager_support.py | 4 +- .../regression_tests/cartpole-sac.yaml | 19 ++ rllib/utils/tf_run_builder.py | 5 +- 20 files changed, 404 insertions(+), 163 deletions(-) create mode 100644 rllib/agents/sac/tests/test_sac.py create mode 100644 rllib/tuned_examples/regression_tests/cartpole-sac.yaml 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):