[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 <ekhliang@gmail.com> * WIP. * Update rllib/agents/trainer.py Co-Authored-By: Eric Liang <ekhliang@gmail.com> * 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 <ekhliang@gmail.com> * 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 <kristian.hartikainen@gmail.com> * Update rllib/agents/sac/sac_policy.py Co-Authored-By: Kristian Hartikainen <kristian.hartikainen@gmail.com> * Update rllib/agents/sac/sac_policy.py Co-Authored-By: Kristian Hartikainen <kristian.hartikainen@gmail.com> * Update rllib/models/tf/tf_action_dist.py Co-Authored-By: Kristian Hartikainen <kristian.hartikainen@gmail.com> * Update rllib/models/tf/tf_action_dist.py Co-Authored-By: Kristian Hartikainen <kristian.hartikainen@gmail.com> * 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 <ekhliang@gmail.com> Co-authored-by: Kristian Hartikainen <kristian.hartikainen@gmail.com>
2026-09-12 12:51:15 +08:00 · 2020-03-06 10:37:12 -08:00
parent 7a33a6ea3c
commit 510c850651
20 changed files with 404 additions and 163 deletions
@@ -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.
@@ -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`_
@@ -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",
 ]

@@ -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",
@@ -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.

@@ -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.
@@ -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
@@ -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,
+)
@@ -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."""
@@ -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),
@@ -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__]))
@@ -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
@@ -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.

@@ -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):
@@ -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:
@@ -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.
@@ -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)

@@ -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):
@@ -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
@@ -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):