[RLlib] DDPG re-factor to fit into RLlib's functional algorithm builder API. (#7934)

rllib/BUILD

@@ -943,6 +943,13 @@ py_test(
     srcs = ["utils/exploration/tests/test_explorations.py"]
 )
 
+py_test(
+    name = "test_parameter_noise",
+    tags = ["utils"],
+    size = "small",
+    srcs = ["utils/exploration/tests/test_parameter_noise.py"]
+)
+
 # Schedules
 py_test(
     name = "test_schedules",

rllib/agents/ddpg/ddpg.py

@@ -1,3 +1,5 @@
+import logging
+
 from ray.rllib.agents.trainer import with_common_config
 from ray.rllib.agents.dqn.dqn import GenericOffPolicyTrainer
 from ray.rllib.agents.ddpg.ddpg_policy import DDPGTFPolicy
@@ -6,6 +8,8 @@ from ray.rllib.utils.deprecation import deprecation_warning, \
 from ray.rllib.utils.exploration.per_worker_ornstein_uhlenbeck_noise import \
     PerWorkerOrnsteinUhlenbeckNoise
 
+logger = logging.getLogger(__name__)
+
 # yapf: disable
 # __sphinx_doc_begin__
 DEFAULT_CONFIG = with_common_config({
@@ -80,12 +84,6 @@ DEFAULT_CONFIG = with_common_config({
     },
     # Number of env steps to optimize for before returning
     "timesteps_per_iteration": 1000,
-
-    # TODO(sven): Move to Exploration API's (ParameterNoise class).
-    # If True parameter space noise will be used for exploration
-    # See https://blog.openai.com/better-exploration-with-parameter-noise/
-    "parameter_noise": False,
-
     # Extra configuration that disables exploration.
     "evaluation_config": {
         "explore": False
@@ -146,6 +144,9 @@ DEFAULT_CONFIG = with_common_config({
     "worker_side_prioritization": False,
     # Prevent iterations from going lower than this time span
     "min_iter_time_s": 1,
+
+    # Deprecated keys.
+    "parameter_noise": DEPRECATED_VALUE,
 })
 # __sphinx_doc_end__
 # yapf: enable
@@ -188,6 +189,18 @@ def validate_config(config):
             config["exploration_config"]["type"] = \
                 PerWorkerOrnsteinUhlenbeckNoise
 
+    if config.get("parameter_noise", DEPRECATED_VALUE) != DEPRECATED_VALUE:
+        deprecation_warning("parameter_noise", "exploration_config={"
+                            "type=ParameterNoise"
+                            "}")
+
+    if config["exploration_config"]["type"] == "ParameterNoise":
+        if config["batch_mode"] != "complete_episodes":
+            logger.warning(
+                "ParameterNoise Exploration requires `batch_mode` to be "
+                "'complete_episodes'. Setting batch_mode=complete_episodes.")
+            config["batch_mode"] = "complete_episodes"
+
 
 DDPGTrainer = GenericOffPolicyTrainer.with_updates(
     name="DDPG",

rllib/agents/ddpg/ddpg_model.py

@@ -0,0 +1,185 @@
+from ray.rllib.models.tf.tf_modelv2 import TFModelV2
+from ray.rllib.utils import try_import_tf
+
+tf = try_import_tf()
+
+
+class DDPGModel(TFModelV2):
+    """Extension of standard TFModel to provide DDPG action- and q-outputs.
+
+    Data flow:
+        obs -> forward() -> model_out
+        model_out -> get_policy_output() -> deterministic actions
+        model_out, actions -> get_q_values() -> Q(s, a)
+        model_out, actions -> get_twin_q_values() -> Q_twin(s, a)
+
+    Note that this class by itself is not a valid model unless you
+    implement forward() in a subclass."""
+
+    def __init__(
+            self,
+            obs_space,
+            action_space,
+            num_outputs,
+            model_config,
+            name,
+            # Extra DDPGActionModel args:
+            actor_hiddens=(256, 256),
+            actor_hidden_activation="relu",
+            critic_hiddens=(256, 256),
+            critic_hidden_activation="relu",
+            twin_q=False,
+            add_layer_norm=False):
+        """Initialize variables of this model.
+
+        Extra model kwargs:
+            actor_hiddens (list): Defines size of hidden layers for the DDPG
+                policy head.
+                These will be used to postprocess the model output for the
+                purposes of computing deterministic actions.
+
+        Note that the core layers for forward() are not defined here, this
+        only defines the layers for the DDPG head. Those layers for forward()
+        should be defined in subclasses of DDPGActionModel.
+        """
+
+        super(DDPGModel, self).__init__(obs_space, action_space, num_outputs,
+                                        model_config, name)
+
+        actor_hidden_activation = getattr(tf.nn, actor_hidden_activation,
+                                          tf.nn.relu)
+        critic_hidden_activation = getattr(tf.nn, critic_hidden_activation,
+                                           tf.nn.relu)
+
+        self.model_out = tf.keras.layers.Input(
+            shape=(num_outputs, ), name="model_out")
+        self.action_dim = action_space.shape[0]
+
+        if actor_hiddens:
+            last_layer = self.model_out
+            for i, n in enumerate(actor_hiddens):
+                last_layer = tf.keras.layers.Dense(
+                    n,
+                    name="actor_hidden_{}".format(i),
+                    activation=actor_hidden_activation)(last_layer)
+                if add_layer_norm:
+                    last_layer = tf.keras.layers.LayerNormalization(
+                        name="LayerNorm_{}".format(i))(last_layer)
+            actor_out = tf.keras.layers.Dense(
+                self.action_dim, activation=None, name="actor_out")(last_layer)
+        else:
+            actor_out = self.model_out
+
+        # Use sigmoid to scale to [0,1], but also double magnitude of input to
+        # emulate behaviour of tanh activation used in DDPG and TD3 papers.
+        def lambda_(x):
+            sigmoid_out = tf.nn.sigmoid(2 * x)
+            # Rescale to actual env policy scale
+            # (shape of sigmoid_out is [batch_size, dim_actions], so we reshape
+            # to get same dims)
+            action_range = (action_space.high - action_space.low)[None]
+            low_action = action_space.low[None]
+            actions = action_range * sigmoid_out + low_action
+            return actions
+
+        actor_out = tf.keras.layers.Lambda(lambda_)(actor_out)
+
+        self.action_model = tf.keras.Model(self.model_out, actor_out)
+        self.register_variables(self.action_model.variables)
+
+        # Build the Q-model(s).
+        self.actions_input = tf.keras.layers.Input(
+            shape=(self.action_dim, ), name="actions")
+
+        def build_q_net(name, observations, actions):
+            # For continuous actions: Feed obs and actions (concatenated)
+            # through the NN.
+            q_net = tf.keras.Sequential([
+                tf.keras.layers.Concatenate(axis=1),
+            ] + [
+                tf.keras.layers.Dense(
+                    units=units,
+                    activation=critic_hidden_activation,
+                    name="{}_hidden_{}".format(name, i))
+                for i, units in enumerate(critic_hiddens)
+            ] + [
+                tf.keras.layers.Dense(
+                    units=1, activation=None, name="{}_out".format(name))
+            ])
+
+            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)
+        self.register_variables(self.q_net.variables)
+
+        if twin_q:
+            self.twin_q_net = build_q_net("twin_q", self.model_out,
+                                          self.actions_input)
+            self.register_variables(self.twin_q_net.variables)
+        else:
+            self.twin_q_net = None
+
+    def get_q_values(self, model_out, actions):
+        """Return the Q estimates for the most recent forward pass.
+
+        This implements Q(s, a).
+
+        Arguments:
+            model_out (Tensor): obs embeddings from the model layers, of shape
+                [BATCH_SIZE, num_outputs].
+            actions (Tensor): Actions to return the Q-values for.
+                Shape: [BATCH_SIZE, action_dim].
+
+        Returns:
+            tensor of shape [BATCH_SIZE].
+        """
+        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):
+        """Same as get_q_values but using the twin Q net.
+
+        This implements the twin Q(s, a).
+
+        Arguments:
+            model_out (Tensor): obs embeddings from the model layers, of shape
+                [BATCH_SIZE, num_outputs].
+            actions (Tensor): Actions to return the Q-values for.
+                Shape: [BATCH_SIZE, action_dim].
+
+        Returns:
+            tensor of shape [BATCH_SIZE].
+        """
+        if actions is not None:
+            return self.twin_q_net([model_out, actions])
+        else:
+            return self.twin_q_net(model_out)
+
+    def get_policy_output(self, model_out):
+        """Return the action output for the most recent forward pass.
+
+        This outputs the support for pi(s). For continuous action spaces, this
+        is the action directly.
+
+        Arguments:
+            model_out (Tensor): obs embeddings from the model layers, of shape
+                [BATCH_SIZE, num_outputs].
+
+        Returns:
+            tensor of shape [BATCH_SIZE, action_out_size]
+        """
+        return self.action_model(model_out)
+
+    def policy_variables(self):
+        """Return the list of variables for the policy net."""
+        return list(self.action_model.variables)
+
+    def q_variables(self):
+        """Return the list of variables for Q / twin Q nets."""
+
+        return self.q_net.variables + (self.twin_q_net.variables
+                                       if self.twin_q_net else [])

rllib/agents/ddpg/ddpg_policy.py

@@ -1,22 +1,28 @@
 from gym.spaces import Box
+import logging
 import numpy as np
 
 import ray
 import ray.experimental.tf_utils
-from ray.rllib.agents.dqn.dqn_tf_policy import postprocess_nstep_and_prio
+from ray.rllib.agents.ddpg.ddpg_model import DDPGModel
+from ray.rllib.agents.ddpg.noop_model import NoopModel
+from ray.rllib.agents.dqn.dqn_tf_policy import postprocess_nstep_and_prio, \
+    PRIO_WEIGHTS
 from ray.rllib.policy.sample_batch import SampleBatch
-from ray.rllib.evaluation.metrics import LEARNER_STATS_KEY
 from ray.rllib.models import ModelCatalog
 from ray.rllib.models.tf.tf_action_dist import Deterministic
 from ray.rllib.utils.annotations import override
 from ray.rllib.utils.error import UnsupportedSpaceException
-from ray.rllib.policy.policy import Policy
 from ray.rllib.policy.tf_policy import TFPolicy
+from ray.rllib.policy.tf_policy_template import build_tf_policy
 from ray.rllib.utils import try_import_tf
-from ray.rllib.utils.tf_ops import huber_loss, minimize_and_clip, scope_vars
+from ray.rllib.utils.tf_ops import huber_loss, minimize_and_clip, \
+    make_tf_callable
 
 tf = try_import_tf()
 
+logger = logging.getLogger(__name__)
+
 ACTION_SCOPE = "action"
 POLICY_SCOPE = "policy"
 POLICY_TARGET_SCOPE = "target_policy"
@@ -25,519 +31,389 @@ Q_TARGET_SCOPE = "target_critic"
 TWIN_Q_SCOPE = "twin_critic"
 TWIN_Q_TARGET_SCOPE = "twin_target_critic"
 
-# Importance sampling weights for prioritized replay
-PRIO_WEIGHTS = "weights"
+
+def build_ddpg_models(policy, observation_space, action_space, config):
+    if config["model"]["custom_model"]:
+        logger.warning(
+            "Setting use_state_preprocessor=True since a custom model "
+            "was specified.")
+        config["use_state_preprocessor"] = True
+
+    if not isinstance(action_space, Box):
+        raise UnsupportedSpaceException(
+            "Action space {} is not supported for DDPG.".format(action_space))
+    elif len(action_space.shape) > 1:
+        raise UnsupportedSpaceException(
+            "Action space has multiple dimensions "
+            "{}. ".format(action_space.shape) +
+            "Consider reshaping this into a single dimension, "
+            "using a Tuple action space, or the multi-agent API.")
+
+    if policy.config["use_state_preprocessor"]:
+        default_model = None  # catalog decides
+        num_outputs = 256  # arbitrary
+        config["model"]["no_final_linear"] = True
+    else:
+        default_model = NoopModel
+        num_outputs = int(np.product(observation_space.shape))
+
+    policy.model = ModelCatalog.get_model_v2(
+        obs_space=observation_space,
+        action_space=action_space,
+        num_outputs=num_outputs,
+        model_config=config["model"],
+        framework="tf",
+        model_interface=DDPGModel,
+        default_model=default_model,
+        name="ddpg_model",
+        actor_hidden_activation=config["actor_hidden_activation"],
+        actor_hiddens=config["actor_hiddens"],
+        critic_hidden_activation=config["critic_hidden_activation"],
+        critic_hiddens=config["critic_hiddens"],
+        twin_q=config["twin_q"],
+        add_layer_norm=(policy.config["exploration_config"].get("type") ==
+                        "ParameterNoise"),
+    )
+
+    policy.target_model = ModelCatalog.get_model_v2(
+        obs_space=observation_space,
+        action_space=action_space,
+        num_outputs=num_outputs,
+        model_config=config["model"],
+        framework="tf",
+        model_interface=DDPGModel,
+        default_model=default_model,
+        name="target_ddpg_model",
+        actor_hidden_activation=config["actor_hidden_activation"],
+        actor_hiddens=config["actor_hiddens"],
+        critic_hidden_activation=config["critic_hidden_activation"],
+        critic_hiddens=config["critic_hiddens"],
+        twin_q=config["twin_q"],
+        add_layer_norm=(policy.config["exploration_config"].get("type") ==
+                        "ParameterNoise"),
+    )
+
+    return policy.model
 
 
-class DDPGPostprocessing:
-    """Implements n-step learning and param noise adjustments."""
+def get_distribution_inputs_and_class(policy,
+                                      model,
+                                      obs_batch,
+                                      *,
+                                      explore=True,
+                                      **kwargs):
+    model_out, _ = model({
+        "obs": obs_batch,
+        "is_training": policy._get_is_training_placeholder()
+    }, [], None)
+    dist_inputs = model.get_policy_output(model_out)
 
-    @override(Policy)
-    def postprocess_trajectory(self,
-                               sample_batch,
-                               other_agent_batches=None,
-                               episode=None):
-        if self.config["parameter_noise"]:
-            # adjust the sigma of parameter space noise
-            states, noisy_actions = [
-                list(x) for x in sample_batch.columns(
-                    [SampleBatch.CUR_OBS, SampleBatch.ACTIONS])
-            ]
-            self.sess.run(self.remove_parameter_noise_op)
-
-            # TODO(sven): This won't work if exploration != Noise, which is
-            #  probably fine as parameter_noise will soon be its own
-            #  Exploration class.
-            clean_actions, cur_noise_scale = self.sess.run(
-                [self.output_actions,
-                 self.exploration.get_info()],
-                feed_dict={
-                    self.cur_observations: states,
-                    self._is_exploring: False,
-                    self._timestep: self.global_timestep,
-                })
-            distance_in_action_space = np.sqrt(
-                np.mean(np.square(clean_actions - noisy_actions)))
-            self.pi_distance = distance_in_action_space
-            if distance_in_action_space < \
-                    self.config["exploration_config"].get("ou_sigma", 0.2) * \
-                    cur_noise_scale:
-                # multiplying the sampled OU noise by noise scale is
-                # equivalent to multiplying the sigma of OU by noise scale
-                self.parameter_noise_sigma_val *= 1.01
-            else:
-                self.parameter_noise_sigma_val /= 1.01
-            self.parameter_noise_sigma.load(
-                self.parameter_noise_sigma_val, session=self.sess)
-
-        return postprocess_nstep_and_prio(self, sample_batch)
+    return dist_inputs, Deterministic, []  # []=state out
 
 
-class DDPGTFPolicy(DDPGPostprocessing, TFPolicy):
-    def __init__(self, observation_space, action_space, config):
-        self.observation_space = observation_space
-        self.action_space = action_space
-        config = dict(ray.rllib.agents.ddpg.ddpg.DEFAULT_CONFIG, **config)
-        if not isinstance(action_space, Box):
-            raise UnsupportedSpaceException(
-                "Action space {} is not supported for DDPG.".format(
-                    action_space))
-        if len(action_space.shape) > 1:
-            raise UnsupportedSpaceException(
-                "Action space has multiple dimensions "
-                "{}. ".format(action_space.shape) +
-                "Consider reshaping this into a single dimension, "
-                "using a Tuple action space, or the multi-agent API.")
+def ddpg_actor_critic_loss(policy, model, _, train_batch):
+    twin_q = policy.config["twin_q"]
+    gamma = policy.config["gamma"]
+    n_step = policy.config["n_step"]
+    use_huber = policy.config["use_huber"]
+    huber_threshold = policy.config["huber_threshold"]
+    l2_reg = policy.config["l2_reg"]
 
-        self.config = config
+    input_dict = {
+        "obs": train_batch[SampleBatch.CUR_OBS],
+        "is_training": policy._get_is_training_placeholder(),
+    }
+    input_dict_next = {
+        "obs": train_batch[SampleBatch.NEXT_OBS],
+        "is_training": policy._get_is_training_placeholder(),
+    }
 
-        # Create global step for counting the number of update operations.
-        self.global_step = tf.train.get_or_create_global_step()
-        # Create sampling timestep placeholder.
-        timestep = tf.placeholder(tf.int32, (), name="timestep")
+    model_out_t, _ = model(input_dict, [], None)
+    model_out_tp1, _ = model(input_dict_next, [], None)
+    target_model_out_tp1, _ = policy.target_model(input_dict_next, [], None)
 
-        # use separate optimizers for actor & critic
-        self._actor_optimizer = tf.train.AdamOptimizer(
-            learning_rate=self.config["actor_lr"])
-        self._critic_optimizer = tf.train.AdamOptimizer(
-            learning_rate=self.config["critic_lr"])
+    # Policy network evaluation.
+    with tf.variable_scope(POLICY_SCOPE, reuse=True):
+        # prev_update_ops = set(tf.get_collection(tf.GraphKeys.UPDATE_OPS))
+        policy_t = model.get_policy_output(model_out_t)
+        # policy_batchnorm_update_ops = list(
+        #    set(tf.get_collection(tf.GraphKeys.UPDATE_OPS)) - prev_update_ops)
 
-        # Observation inputs.
-        self.cur_observations = tf.placeholder(
-            tf.float32,
-            shape=(None, ) + observation_space.shape,
-            name="cur_obs")
+    with tf.variable_scope(POLICY_TARGET_SCOPE):
+        policy_tp1 = \
+            policy.target_model.get_policy_output(target_model_out_tp1)
 
-        with tf.variable_scope(POLICY_SCOPE) as scope:
-            self._distribution_inputs, self.policy_model = \
-                self._build_policy_network(
-                    self.cur_observations, observation_space, action_space)
-            self.policy_vars = scope_vars(scope.name)
-        self.model = self.policy_model
-
-        # Noise vars for P network except for layer normalization vars
-        if self.config["parameter_noise"]:
-            self._build_parameter_noise([
-                var for var in self.policy_vars if "LayerNorm" not in var.name
-            ])
-
-        # Create exploration component.
-        self.exploration = self._create_exploration()
-        explore = tf.placeholder_with_default(True, (), name="is_exploring")
-        # Action outputs.
-        with tf.variable_scope(ACTION_SCOPE):
-            self.output_actions, _ = self.exploration.get_exploration_action(
-                action_distribution=Deterministic(self._distribution_inputs,
-                                                  self.model),
-                timestep=timestep,
-                explore=explore)
-
-        # Replay inputs.
-        self.obs_t = tf.placeholder(
-            tf.float32,
-            shape=(None, ) + observation_space.shape,
-            name="observation")
-        self.act_t = tf.placeholder(
-            tf.float32, shape=(None, ) + action_space.shape, name="action")
-        self.rew_t = tf.placeholder(tf.float32, [None], name="reward")
-        self.obs_tp1 = tf.placeholder(
-            tf.float32, shape=(None, ) + observation_space.shape)
-        self.done_mask = tf.placeholder(tf.float32, [None], name="done")
-        self.importance_weights = tf.placeholder(
-            tf.float32, [None], name="weight")
-
-        # policy network evaluation
-        with tf.variable_scope(POLICY_SCOPE, reuse=True) as scope:
-            prev_update_ops = set(tf.get_collection(tf.GraphKeys.UPDATE_OPS))
-            self.policy_t, _ = self._build_policy_network(
-                self.obs_t, observation_space, action_space)
-            policy_batchnorm_update_ops = list(
-                set(tf.get_collection(tf.GraphKeys.UPDATE_OPS)) -
-                prev_update_ops)
-
-        # target policy network evaluation
-        with tf.variable_scope(POLICY_TARGET_SCOPE) as scope:
-            policy_tp1, _ = self._build_policy_network(
-                self.obs_tp1, observation_space, action_space)
-            target_policy_vars = scope_vars(scope.name)
-
-        # Action outputs
-        with tf.variable_scope(ACTION_SCOPE, reuse=True):
-            if config["smooth_target_policy"]:
-                target_noise_clip = self.config["target_noise_clip"]
-                clipped_normal_sample = tf.clip_by_value(
-                    tf.random_normal(
-                        tf.shape(policy_tp1),
-                        stddev=self.config["target_noise"]),
-                    -target_noise_clip, target_noise_clip)
-                policy_tp1_smoothed = tf.clip_by_value(
-                    policy_tp1 + clipped_normal_sample,
-                    action_space.low * tf.ones_like(policy_tp1),
-                    action_space.high * tf.ones_like(policy_tp1))
-            else:
-                # no smoothing, just use deterministic actions
-                policy_tp1_smoothed = policy_tp1
-
-        # q network evaluation
-        prev_update_ops = set(tf.get_collection(tf.GraphKeys.UPDATE_OPS))
-        with tf.variable_scope(Q_SCOPE) as scope:
-            # Q-values for given actions & observations in given current
-            q_t, self.q_model = self._build_q_network(
-                self.obs_t, observation_space, action_space, self.act_t)
-            self.q_func_vars = scope_vars(scope.name)
-        self.stats = {
-            "mean_q": tf.reduce_mean(q_t),
-            "max_q": tf.reduce_max(q_t),
-            "min_q": tf.reduce_min(q_t),
-        }
-        with tf.variable_scope(Q_SCOPE, reuse=True):
-            # Q-values for current policy (no noise) in given current state
-            q_t_det_policy, _ = self._build_q_network(
-                self.obs_t, observation_space, action_space, self.policy_t)
-        if self.config["twin_q"]:
-            with tf.variable_scope(TWIN_Q_SCOPE) as scope:
-                twin_q_t, self.twin_q_model = self._build_q_network(
-                    self.obs_t, observation_space, action_space, self.act_t)
-                self.twin_q_func_vars = scope_vars(scope.name)
-        q_batchnorm_update_ops = list(
-            set(tf.get_collection(tf.GraphKeys.UPDATE_OPS)) - prev_update_ops)
-
-        # target q network evaluation
-        with tf.variable_scope(Q_TARGET_SCOPE) as scope:
-            q_tp1, _ = self._build_q_network(self.obs_tp1, observation_space,
-                                             action_space, policy_tp1_smoothed)
-            target_q_func_vars = scope_vars(scope.name)
-        if self.config["twin_q"]:
-            with tf.variable_scope(TWIN_Q_TARGET_SCOPE) as scope:
-                twin_q_tp1, _ = self._build_q_network(
-                    self.obs_tp1, observation_space, action_space,
-                    policy_tp1_smoothed)
-                twin_target_q_func_vars = scope_vars(scope.name)
-
-        if self.config["twin_q"]:
-            self.critic_loss, self.actor_loss, self.td_error \
-                = self._build_actor_critic_loss(
-                    q_t, q_tp1, q_t_det_policy, twin_q_t=twin_q_t,
-                    twin_q_tp1=twin_q_tp1)
+    # Action outputs.
+    with tf.variable_scope(ACTION_SCOPE, reuse=True):
+        if policy.config["smooth_target_policy"]:
+            target_noise_clip = policy.config["target_noise_clip"]
+            clipped_normal_sample = tf.clip_by_value(
+                tf.random_normal(
+                    tf.shape(policy_tp1),
+                    stddev=policy.config["target_noise"]), -target_noise_clip,
+                target_noise_clip)
+            policy_tp1_smoothed = tf.clip_by_value(
+                policy_tp1 + clipped_normal_sample,
+                policy.action_space.low * tf.ones_like(policy_tp1),
+                policy.action_space.high * tf.ones_like(policy_tp1))
         else:
-            self.critic_loss, self.actor_loss, self.td_error \
-                = self._build_actor_critic_loss(
-                    q_t, q_tp1, q_t_det_policy)
+            # No smoothing, just use deterministic actions.
+            policy_tp1_smoothed = policy_tp1
 
-        if config["l2_reg"] is not None:
-            for var in self.policy_vars:
-                if "bias" not in var.name:
-                    self.actor_loss += (config["l2_reg"] * tf.nn.l2_loss(var))
-            for var in self.q_func_vars:
-                if "bias" not in var.name:
-                    self.critic_loss += (config["l2_reg"] * tf.nn.l2_loss(var))
-            if self.config["twin_q"]:
-                for var in self.twin_q_func_vars:
-                    if "bias" not in var.name:
-                        self.critic_loss += (
-                            config["l2_reg"] * tf.nn.l2_loss(var))
+    # Q-net(s) evaluation.
+    # prev_update_ops = set(tf.get_collection(tf.GraphKeys.UPDATE_OPS))
+    with tf.variable_scope(Q_SCOPE):
+        # Q-values for given actions & observations in given current
+        q_t = model.get_q_values(model_out_t, train_batch[SampleBatch.ACTIONS])
 
-        # update_target_fn will be called periodically to copy Q network to
-        # target Q network
-        self.tau_value = config.get("tau")
-        self.tau = tf.placeholder(tf.float32, (), name="tau")
-        update_target_expr = []
-        for var, var_target in zip(
-                sorted(self.q_func_vars, key=lambda v: v.name),
-                sorted(target_q_func_vars, key=lambda v: v.name)):
-            update_target_expr.append(
-                var_target.assign(self.tau * var +
-                                  (1.0 - self.tau) * var_target))
-        if self.config["twin_q"]:
-            for var, var_target in zip(
-                    sorted(self.twin_q_func_vars, key=lambda v: v.name),
-                    sorted(twin_target_q_func_vars, key=lambda v: v.name)):
+    with tf.variable_scope(Q_SCOPE, reuse=True):
+        # 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 twin_q:
+        with tf.variable_scope(TWIN_Q_SCOPE):
+            twin_q_t = model.get_twin_q_values(
+                model_out_t, train_batch[SampleBatch.ACTIONS])
+    # q_batchnorm_update_ops = list(
+    #     set(tf.get_collection(tf.GraphKeys.UPDATE_OPS)) - prev_update_ops)
+
+    # Target q-net(s) evaluation.
+    with tf.variable_scope(Q_TARGET_SCOPE):
+        q_tp1 = policy.target_model.get_q_values(target_model_out_tp1,
+                                                 policy_tp1_smoothed)
+
+    if twin_q:
+        with tf.variable_scope(TWIN_Q_TARGET_SCOPE):
+            twin_q_tp1 = policy.target_model.get_twin_q_values(
+                target_model_out_tp1, policy_tp1_smoothed)
+
+    q_t_selected = tf.squeeze(q_t, axis=len(q_t.shape) - 1)
+    if twin_q:
+        twin_q_t_selected = tf.squeeze(twin_q_t, axis=len(q_t.shape) - 1)
+        q_tp1 = tf.minimum(q_tp1, twin_q_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
+
+    # Compute RHS of bellman equation.
+    q_t_selected_target = tf.stop_gradient(train_batch[SampleBatch.REWARDS] +
+                                           gamma**n_step * q_tp1_best_masked)
+
+    # Compute the error (potentially clipped).
+    if twin_q:
+        td_error = q_t_selected - q_t_selected_target
+        twin_td_error = twin_q_t_selected - q_t_selected_target
+        td_error = td_error + twin_td_error
+        if use_huber:
+            errors = huber_loss(td_error, huber_threshold) \
+                + huber_loss(twin_td_error, huber_threshold)
+        else:
+            errors = 0.5 * tf.square(td_error) + 0.5 * tf.square(twin_td_error)
+    else:
+        td_error = q_t_selected - q_t_selected_target
+        if use_huber:
+            errors = huber_loss(td_error, huber_threshold)
+        else:
+            errors = 0.5 * tf.square(td_error)
+
+    critic_loss = tf.reduce_mean(train_batch[PRIO_WEIGHTS] * errors)
+    actor_loss = -tf.reduce_mean(q_t_det_policy)
+
+    # Add l2-regularization if required.
+    if l2_reg is not None:
+        for var in policy.model.policy_variables():
+            if "bias" not in var.name:
+                actor_loss += (l2_reg * tf.nn.l2_loss(var))
+        for var in policy.model.q_variables():
+            if "bias" not in var.name:
+                critic_loss += (l2_reg * tf.nn.l2_loss(var))
+
+    # Model self-supervised losses.
+    if policy.config["use_state_preprocessor"]:
+        # Expand input_dict in case custom_loss' need them.
+        input_dict[SampleBatch.ACTIONS] = train_batch[SampleBatch.ACTIONS]
+        input_dict[SampleBatch.REWARDS] = train_batch[SampleBatch.REWARDS]
+        input_dict[SampleBatch.DONES] = train_batch[SampleBatch.DONES]
+        input_dict[SampleBatch.NEXT_OBS] = train_batch[SampleBatch.NEXT_OBS]
+        actor_loss, critic_loss = model.custom_loss([actor_loss, critic_loss],
+                                                    input_dict)
+
+    # Store values for stats function.
+    policy.actor_loss = actor_loss
+    policy.critic_loss = critic_loss
+    policy.td_error = td_error
+    policy.q_t = q_t
+
+    # Return one loss value (even though we treat them separately in our
+    # 2 optimizers: actor and critic).
+    return policy.critic_loss + policy.actor_loss
+
+
+def make_ddpg_optimizers(policy, config):
+    # Create separate optimizers for actor & critic losses.
+    policy._actor_optimizer = tf.train.AdamOptimizer(
+        learning_rate=config["actor_lr"])
+    policy._critic_optimizer = tf.train.AdamOptimizer(
+        learning_rate=config["critic_lr"])
+    return None
+
+    # TFPolicy.__init__(
+    #    self,
+    #    observation_space,
+    #    action_space,
+    #    self.config,
+    #    self.sess,
+    #    #obs_input=self.cur_observations,
+    #    sampled_action=self.output_actions,
+    #    loss=self.actor_loss + self.critic_loss,
+    #    loss_inputs=self.loss_inputs,
+    #    update_ops=q_batchnorm_update_ops + policy_batchnorm_update_ops,
+    #    explore=explore,
+    #    dist_inputs=self._distribution_inputs,
+    #    dist_class=Deterministic,
+    #    timestep=timestep)
+
+
+def build_apply_op(policy, optimizer, grads_and_vars):
+    # For policy gradient, update policy net one time v.s.
+    # update critic net `policy_delay` time(s).
+    should_apply_actor_opt = tf.equal(
+        tf.mod(policy.global_step, policy.config["policy_delay"]), 0)
+
+    def make_apply_op():
+        return policy._actor_optimizer.apply_gradients(
+            policy._actor_grads_and_vars)
+
+    actor_op = tf.cond(
+        should_apply_actor_opt,
+        true_fn=make_apply_op,
+        false_fn=lambda: tf.no_op())
+    critic_op = policy._critic_optimizer.apply_gradients(
+        policy._critic_grads_and_vars)
+    # Increment global step & apply ops.
+    with tf.control_dependencies([tf.assign_add(policy.global_step, 1)]):
+        return tf.group(actor_op, critic_op)
+
+
+def gradients_fn(policy, optimizer, loss):
+    if policy.config["grad_norm_clipping"] is not None:
+        actor_grads_and_vars = minimize_and_clip(
+            policy._actor_optimizer,
+            policy.actor_loss,
+            var_list=policy.model.policy_variables(),
+            clip_val=policy.config["grad_norm_clipping"])
+        critic_grads_and_vars = minimize_and_clip(
+            policy._critic_optimizer,
+            policy.critic_loss,
+            var_list=policy.model.q_variables(),
+            clip_val=policy.config["grad_norm_clipping"])
+    else:
+        actor_grads_and_vars = policy._actor_optimizer.compute_gradients(
+            policy.actor_loss, var_list=policy.model.policy_variables())
+        critic_grads_and_vars = policy._critic_optimizer.compute_gradients(
+            policy.critic_loss, var_list=policy.model.q_variables())
+    # Save these for later use in build_apply_op.
+    policy._actor_grads_and_vars = [(g, v) for (g, v) in actor_grads_and_vars
+                                    if g is not None]
+    policy._critic_grads_and_vars = [(g, v) for (g, v) in critic_grads_and_vars
+                                     if g is not None]
+    grads_and_vars = policy._actor_grads_and_vars + \
+        policy._critic_grads_and_vars
+    return grads_and_vars
+
+
+def build_ddpg_stats(policy, batch):
+    stats = {
+        "mean_q": tf.reduce_mean(policy.q_t),
+        "max_q": tf.reduce_max(policy.q_t),
+        "min_q": tf.reduce_min(policy.q_t),
+    }
+    return stats
+
+
+def before_init_fn(policy, obs_space, action_space, config):
+    # Create global step for counting the number of update operations.
+    policy.global_step = tf.train.get_or_create_global_step()
+
+
+class ComputeTDErrorMixin:
+    def __init__(self, loss_fn):
+        @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 errors attribute
+            # (one TD-error value per item in batch to update PR weights).
+            loss_fn(
+                self, self.model, None, {
+                    SampleBatch.CUR_OBS: tf.convert_to_tensor(obs_t),
+                    SampleBatch.ACTIONS: tf.convert_to_tensor(act_t),
+                    SampleBatch.REWARDS: tf.convert_to_tensor(rew_t),
+                    SampleBatch.NEXT_OBS: tf.convert_to_tensor(obs_tp1),
+                    SampleBatch.DONES: tf.convert_to_tensor(done_mask),
+                    PRIO_WEIGHTS: tf.convert_to_tensor(importance_weights),
+                })
+            # `self.td_error` is set in loss_fn.
+            return self.td_error
+
+        self.compute_td_error = compute_td_error
+
+
+def setup_mid_mixins(policy, obs_space, action_space, config):
+    ComputeTDErrorMixin.__init__(policy, ddpg_actor_critic_loss)
+
+
+class TargetNetworkMixin:
+    def __init__(self, config):
+        @make_tf_callable(self.get_session())
+        def update_target_fn(tau):
+            tau = tf.convert_to_tensor(tau, dtype=tf.float32)
+            update_target_expr = []
+            model_vars = self.model.trainable_variables()
+            target_model_vars = self.target_model.trainable_variables()
+            assert len(model_vars) == len(target_model_vars), \
+                (model_vars, target_model_vars)
+            for var, var_target in zip(model_vars, target_model_vars):
                 update_target_expr.append(
-                    var_target.assign(self.tau * var +
-                                      (1.0 - self.tau) * var_target))
-        for var, var_target in zip(
-                sorted(self.policy_vars, key=lambda v: v.name),
-                sorted(target_policy_vars, key=lambda v: v.name)):
-            update_target_expr.append(
-                var_target.assign(self.tau * var +
-                                  (1.0 - self.tau) * var_target))
-        self.update_target_expr = tf.group(*update_target_expr)
+                    var_target.assign(tau * var + (1.0 - tau) * var_target))
+                logger.debug("Update target op {}".format(var_target))
+            return tf.group(*update_target_expr)
 
-        self.sess = tf.get_default_session()
-        self.loss_inputs = [
-            (SampleBatch.CUR_OBS, self.obs_t),
-            (SampleBatch.ACTIONS, self.act_t),
-            (SampleBatch.REWARDS, self.rew_t),
-            (SampleBatch.NEXT_OBS, self.obs_tp1),
-            (SampleBatch.DONES, self.done_mask),
-            (PRIO_WEIGHTS, self.importance_weights),
-        ]
-        input_dict = dict(self.loss_inputs)
-
-        if self.config["use_state_preprocessor"]:
-            # Model self-supervised losses
-            self.actor_loss = self.policy_model.custom_loss(
-                self.actor_loss, input_dict)
-            self.critic_loss = self.q_model.custom_loss(
-                self.critic_loss, input_dict)
-            if self.config["twin_q"]:
-                self.critic_loss = self.twin_q_model.custom_loss(
-                    self.critic_loss, input_dict)
-
-        TFPolicy.__init__(
-            self,
-            observation_space,
-            action_space,
-            self.config,
-            self.sess,
-            obs_input=self.cur_observations,
-            sampled_action=self.output_actions,
-            loss=self.actor_loss + self.critic_loss,
-            loss_inputs=self.loss_inputs,
-            update_ops=q_batchnorm_update_ops + policy_batchnorm_update_ops,
-            explore=explore,
-            dist_inputs=self._distribution_inputs,
-            dist_class=Deterministic,
-            timestep=timestep)
-        self.sess.run(tf.global_variables_initializer())
-
-        # Note that this encompasses both the policy and Q-value networks and
-        # their corresponding target networks
-        self.variables = ray.experimental.tf_utils.TensorFlowVariables(
-            tf.group(q_t_det_policy, q_tp1, self._actor_optimizer.variables(),
-                     self._critic_optimizer.variables()), self.sess)
-
-        # Hard initial update
+        # Hard initial update.
+        self._do_update = update_target_fn
         self.update_target(tau=1.0)
 
-    @override(TFPolicy)
-    def optimizer(self):
-        # we don't use this because we have two separate optimisers
-        return None
-
-    @override(TFPolicy)
-    def build_apply_op(self, optimizer, grads_and_vars):
-        # for policy gradient, update policy net one time v.s.
-        # update critic net `policy_delay` time(s)
-        should_apply_actor_opt = tf.equal(
-            tf.mod(self.global_step, self.config["policy_delay"]), 0)
-
-        def make_apply_op():
-            return self._actor_optimizer.apply_gradients(
-                self._actor_grads_and_vars)
-
-        actor_op = tf.cond(
-            should_apply_actor_opt,
-            true_fn=make_apply_op,
-            false_fn=lambda: tf.no_op())
-        critic_op = self._critic_optimizer.apply_gradients(
-            self._critic_grads_and_vars)
-        # increment global step & apply ops
-        with tf.control_dependencies([tf.assign_add(self.global_step, 1)]):
-            return tf.group(actor_op, critic_op)
-
-    @override(TFPolicy)
-    def gradients(self, optimizer, loss):
-        if self.config["grad_norm_clipping"] is not None:
-            actor_grads_and_vars = minimize_and_clip(
-                self._actor_optimizer,
-                self.actor_loss,
-                var_list=self.policy_vars,
-                clip_val=self.config["grad_norm_clipping"])
-            critic_grads_and_vars = minimize_and_clip(
-                self._critic_optimizer,
-                self.critic_loss,
-                var_list=self.q_func_vars + self.twin_q_func_vars
-                if self.config["twin_q"] else self.q_func_vars,
-                clip_val=self.config["grad_norm_clipping"])
-        else:
-            actor_grads_and_vars = self._actor_optimizer.compute_gradients(
-                self.actor_loss, var_list=self.policy_vars)
-            if self.config["twin_q"]:
-                critic_vars = self.q_func_vars + self.twin_q_func_vars
-            else:
-                critic_vars = self.q_func_vars
-            critic_grads_and_vars = self._critic_optimizer.compute_gradients(
-                self.critic_loss, var_list=critic_vars)
-        # save these for later use in build_apply_op
-        self._actor_grads_and_vars = [(g, v) for (g, v) in actor_grads_and_vars
-                                      if g is not None]
-        self._critic_grads_and_vars = [(g, v)
-                                       for (g, v) in critic_grads_and_vars
-                                       if g is not None]
-        grads_and_vars = self._actor_grads_and_vars \
-            + self._critic_grads_and_vars
-        return grads_and_vars
-
-    @override(TFPolicy)
-    def extra_compute_grad_fetches(self):
-        return {
-            "td_error": self.td_error,
-            LEARNER_STATS_KEY: self.stats,
-        }
-
-    @override(TFPolicy)
-    def get_weights(self):
-        return self.variables.get_weights()
-
-    @override(TFPolicy)
-    def set_weights(self, weights):
-        self.variables.set_weights(weights)
-
-    def _build_q_network(self, obs, obs_space, action_space, actions):
-        if self.config["use_state_preprocessor"]:
-            q_model = ModelCatalog.get_model({
-                "obs": obs,
-                "is_training": self._get_is_training_placeholder(),
-            }, obs_space, action_space, 1, self.config["model"])
-            q_out = tf.concat([q_model.last_layer, actions], axis=1)
-        else:
-            q_model = None
-            q_out = tf.concat([obs, actions], axis=1)
-
-        activation = getattr(tf.nn, self.config["critic_hidden_activation"])
-        for hidden in self.config["critic_hiddens"]:
-            q_out = tf.layers.dense(q_out, units=hidden, activation=activation)
-        q_values = tf.layers.dense(q_out, units=1, activation=None)
-
-        return q_values, q_model
-
-    def _build_policy_network(self, obs, obs_space, action_space):
-        if self.config["use_state_preprocessor"]:
-            model = ModelCatalog.get_model({
-                "obs": obs,
-                "is_training": self._get_is_training_placeholder(),
-            }, obs_space, action_space, 1, self.config["model"])
-            action_out = model.last_layer
-        else:
-            model = None
-            action_out = obs
-
-        activation = getattr(tf.nn, self.config["actor_hidden_activation"])
-        for hidden in self.config["actor_hiddens"]:
-            action_out = tf.layers.dense(
-                action_out, units=hidden, activation=activation)
-            if self.config["parameter_noise"]:
-                action_out = tf.keras.layers.LayerNormalization()(action_out)
-        action_out = tf.layers.dense(
-            action_out, units=action_space.shape[0], activation=None)
-
-        # Use sigmoid to scale to [0,1], but also double magnitude of input to
-        # emulate behaviour of tanh activation used in DDPG and TD3 papers.
-        sigmoid_out = tf.nn.sigmoid(2 * action_out)
-        # Rescale to actual env policy scale
-        # (shape of sigmoid_out is [batch_size, dim_actions], so we reshape to
-        # get same dims)
-        action_range = (action_space.high - action_space.low)[None]
-        low_action = action_space.low[None]
-        actions = action_range * sigmoid_out + low_action
-
-        return actions, model
-
-    def _build_actor_critic_loss(self,
-                                 q_t,
-                                 q_tp1,
-                                 q_t_det_policy,
-                                 twin_q_t=None,
-                                 twin_q_tp1=None):
-        twin_q = self.config["twin_q"]
-        gamma = self.config["gamma"]
-        n_step = self.config["n_step"]
-        use_huber = self.config["use_huber"]
-        huber_threshold = self.config["huber_threshold"]
-
-        q_t_selected = tf.squeeze(q_t, axis=len(q_t.shape) - 1)
-        if twin_q:
-            twin_q_t_selected = tf.squeeze(twin_q_t, axis=len(q_t.shape) - 1)
-            q_tp1 = tf.minimum(q_tp1, twin_q_tp1)
-
-        q_tp1_best = tf.squeeze(input=q_tp1, axis=len(q_tp1.shape) - 1)
-        q_tp1_best_masked = (1.0 - self.done_mask) * q_tp1_best
-
-        # compute RHS of bellman equation
-        q_t_selected_target = tf.stop_gradient(
-            self.rew_t + gamma**n_step * q_tp1_best_masked)
-
-        # compute the error (potentially clipped)
-        if twin_q:
-            td_error = q_t_selected - q_t_selected_target
-            twin_td_error = twin_q_t_selected - q_t_selected_target
-            td_error = td_error + twin_td_error
-            if use_huber:
-                errors = huber_loss(td_error, huber_threshold) \
-                    + huber_loss(twin_td_error, huber_threshold)
-            else:
-                errors = 0.5 * tf.square(td_error) + 0.5 * tf.square(
-                    twin_td_error)
-        else:
-            td_error = q_t_selected - q_t_selected_target
-            if use_huber:
-                errors = huber_loss(td_error, huber_threshold)
-            else:
-                errors = 0.5 * tf.square(td_error)
-
-        critic_loss = tf.reduce_mean(self.importance_weights * errors)
-        actor_loss = -tf.reduce_mean(q_t_det_policy)
-        return critic_loss, actor_loss, td_error
-
-    def _build_parameter_noise(self, pnet_params):
-        self.parameter_noise_sigma_val = \
-            self.config["exploration_config"].get("ou_sigma", 0.2)
-        self.parameter_noise_sigma = tf.get_variable(
-            initializer=tf.constant_initializer(
-                self.parameter_noise_sigma_val),
-            name="parameter_noise_sigma",
-            shape=(),
-            trainable=False,
-            dtype=tf.float32)
-        self.parameter_noise = list()
-        # No need to add any noise on LayerNorm parameters
-        for var in pnet_params:
-            noise_var = tf.get_variable(
-                name=var.name.split(":")[0] + "_noise",
-                shape=var.shape,
-                initializer=tf.constant_initializer(.0),
-                trainable=False)
-            self.parameter_noise.append(noise_var)
-        remove_noise_ops = list()
-        for var, var_noise in zip(pnet_params, self.parameter_noise):
-            remove_noise_ops.append(tf.assign_add(var, -var_noise))
-        self.remove_parameter_noise_op = tf.group(*tuple(remove_noise_ops))
-        generate_noise_ops = list()
-        for var_noise in self.parameter_noise:
-            generate_noise_ops.append(
-                tf.assign(
-                    var_noise,
-                    tf.random_normal(
-                        shape=var_noise.shape,
-                        stddev=self.parameter_noise_sigma)))
-        with tf.control_dependencies(generate_noise_ops):
-            add_noise_ops = list()
-            for var, var_noise in zip(pnet_params, self.parameter_noise):
-                add_noise_ops.append(tf.assign_add(var, var_noise))
-            self.add_noise_op = tf.group(*tuple(add_noise_ops))
-        self.pi_distance = None
-
-    def compute_td_error(self, obs_t, act_t, rew_t, obs_tp1, done_mask,
-                         importance_weights):
-        td_err = self.sess.run(
-            self.td_error,
-            feed_dict={
-                self.obs_t: [np.array(ob) for ob in obs_t],
-                self.act_t: act_t,
-                self.rew_t: rew_t,
-                self.obs_tp1: [np.array(ob) for ob in obs_tp1],
-                self.done_mask: done_mask,
-                self.importance_weights: importance_weights
-            })
-        return td_err
-
-    def add_parameter_noise(self):
-        if self.config["parameter_noise"]:
-            self.sess.run(self.add_noise_op)
-
-    # support both hard and soft sync
+    # Support both hard and soft sync.
     def update_target(self, tau=None):
-        tau = tau or self.tau_value
-        return self.sess.run(
-            self.update_target_expr, feed_dict={self.tau: tau})
+        self._do_update(np.float32(tau or self.config.get("tau")))
+
+    @override(TFPolicy)
+    def variables(self):
+        return self.model.variables() + self.target_model.variables()
+
+
+def setup_late_mixins(policy, obs_space, action_space, config):
+    TargetNetworkMixin.__init__(policy, config)
+
+
+DDPGTFPolicy = build_tf_policy(
+    name="DQNTFPolicy",
+    get_default_config=lambda: ray.rllib.agents.ddpg.ddpg.DEFAULT_CONFIG,
+    make_model=build_ddpg_models,
+    action_distribution_fn=get_distribution_inputs_and_class,
+    loss_fn=ddpg_actor_critic_loss,
+    stats_fn=build_ddpg_stats,
+    postprocess_fn=postprocess_nstep_and_prio,
+    optimizer_fn=make_ddpg_optimizers,
+    gradients_fn=gradients_fn,
+    apply_gradients_fn=build_apply_op,
+    extra_learn_fetches_fn=lambda policy: {"td_error": policy.td_error},
+    before_init=before_init_fn,
+    before_loss_init=setup_mid_mixins,
+    after_init=setup_late_mixins,
+    obs_include_prev_action_reward=False,
+    mixins=[
+        TargetNetworkMixin,
+        ComputeTDErrorMixin,
+    ])

rllib/agents/ddpg/tests/test_ddpg.py

@@ -14,10 +14,11 @@ class TestDDPG(unittest.TestCase):
         config = ddpg.DEFAULT_CONFIG.copy()
         config["num_workers"] = 0  # Run locally.
 
+        num_iterations = 2
+
         # Test against all frameworks.
         for _ in framework_iterator(config, "tf"):
             trainer = ddpg.DDPGTrainer(config=config, env="Pendulum-v0")
-            num_iterations = 2
             for i in range(num_iterations):
                 results = trainer.train()
                 print(results)

rllib/agents/dqn/tests/test_dqn.py

@@ -100,167 +100,6 @@ class TestDQN(unittest.TestCase):
                 actions.append(trainer.compute_action(obs))
             check(np.std(actions), 0.0, false=True)
 
-    def test_dqn_parameter_noise_exploration(self):
-        """Tests, whether a DQN Agent works with ParameterNoise."""
-        obs = np.array(0)
-        core_config = dqn.DEFAULT_CONFIG.copy()
-        core_config["num_workers"] = 0  # Run locally.
-        core_config["env_config"] = {"is_slippery": False, "map_name": "4x4"}
-
-        # Test against all frameworks.
-        for fw in framework_iterator(core_config):
-            config = core_config.copy()
-
-            # DQN with ParameterNoise exploration (config["explore"]=True).
-            # ----
-            config["exploration_config"] = {"type": "ParameterNoise"}
-            config["explore"] = True
-
-            trainer = dqn.DQNTrainer(config=config, env="FrozenLake-v0")
-            policy = trainer.get_policy()
-            p_sess = getattr(policy, "_sess", None)
-            self.assertFalse(policy.exploration.weights_are_currently_noisy)
-            noise_before = self._get_current_noise(policy, fw)
-            check(noise_before, 0.0)
-            initial_weights = self._get_current_weight(policy, fw)
-
-            # Pseudo-start an episode and compare the weights before and after.
-            policy.exploration.on_episode_start(policy, tf_sess=p_sess)
-            self.assertFalse(policy.exploration.weights_are_currently_noisy)
-            noise_after_ep_start = self._get_current_noise(policy, fw)
-            weights_after_ep_start = self._get_current_weight(policy, fw)
-            # Should be the same, as we don't do anything at the beginning of
-            # the episode, only one step later.
-            check(noise_after_ep_start, noise_before)
-            check(initial_weights, weights_after_ep_start)
-
-            # Setting explore=False should always return the same action.
-            a_ = trainer.compute_action(obs, explore=False)
-            self.assertFalse(policy.exploration.weights_are_currently_noisy)
-            noise = self._get_current_noise(policy, fw)
-            # We sampled the first noise (not zero anymore).
-            check(noise, 0.0, false=True)
-            # But still not applied b/c explore=False.
-            check(self._get_current_weight(policy, fw), initial_weights)
-            for _ in range(10):
-                a = trainer.compute_action(obs, explore=False)
-                check(a, a_)
-                # Noise never gets applied.
-                check(self._get_current_weight(policy, fw), initial_weights)
-                self.assertFalse(
-                    policy.exploration.weights_are_currently_noisy)
-
-            # Explore=None (default: True) should return different actions.
-            # However, this is only due to the underlying epsilon-greedy
-            # exploration.
-            actions = []
-            current_weight = None
-            for _ in range(10):
-                actions.append(trainer.compute_action(obs))
-                self.assertTrue(policy.exploration.weights_are_currently_noisy)
-                # Now, noise actually got applied (explore=True).
-                current_weight = self._get_current_weight(policy, fw)
-                check(current_weight, initial_weights, false=True)
-                check(current_weight, initial_weights + noise)
-            check(np.std(actions), 0.0, false=True)
-
-            # Pseudo-end the episode and compare weights again.
-            # Make sure they are the original ones.
-            policy.exploration.on_episode_end(policy, tf_sess=p_sess)
-            weights_after_ep_end = self._get_current_weight(policy, fw)
-            check(current_weight - noise, weights_after_ep_end, decimals=5)
-
-            # DQN with ParameterNoise exploration (config["explore"]=False).
-            # ----
-            config = core_config.copy()
-            config["exploration_config"] = {"type": "ParameterNoise"}
-            config["explore"] = False
-            trainer = dqn.DQNTrainer(config=config, env="FrozenLake-v0")
-            policy = trainer.get_policy()
-            p_sess = getattr(policy, "_sess", None)
-            self.assertFalse(policy.exploration.weights_are_currently_noisy)
-            initial_weights = self._get_current_weight(policy, fw)
-
-            # Noise before anything (should be 0.0, no episode started yet).
-            noise = self._get_current_noise(policy, fw)
-            check(noise, 0.0)
-
-            # Pseudo-start an episode and compare the weights before and after
-            # (they should be the same).
-            policy.exploration.on_episode_start(policy, tf_sess=p_sess)
-            self.assertFalse(policy.exploration.weights_are_currently_noisy)
-
-            # Should be the same, as we don't do anything at the beginning of
-            # the episode, only one step later.
-            noise = self._get_current_noise(policy, fw)
-            check(noise, 0.0)
-            noisy_weights = self._get_current_weight(policy, fw)
-            check(initial_weights, noisy_weights)
-
-            # Setting explore=False or None should always return the same
-            # action.
-            a_ = trainer.compute_action(obs, explore=False)
-            # Now we have re-sampled.
-            noise = self._get_current_noise(policy, fw)
-            check(noise, 0.0, false=True)
-            for _ in range(5):
-                a = trainer.compute_action(obs, explore=None)
-                check(a, a_)
-                a = trainer.compute_action(obs, explore=False)
-                check(a, a_)
-
-            # Pseudo-end the episode and compare weights again.
-            # Make sure they are the original ones (no noise permanently
-            # applied throughout the episode).
-            policy.exploration.on_episode_end(policy, tf_sess=p_sess)
-            weights_after_episode_end = self._get_current_weight(policy, fw)
-            check(initial_weights, weights_after_episode_end)
-            # Noise should still be the same (re-sampling only happens at
-            # beginning of episode).
-            noise_after = self._get_current_noise(policy, fw)
-            check(noise, noise_after)
-
-            # Switch off EpsilonGreedy underlying exploration.
-            # ----
-            config = core_config.copy()
-            config["exploration_config"] = {
-                "type": "ParameterNoise",
-                "sub_exploration": {
-                    "type": "EpsilonGreedy",
-                    "action_space": trainer.get_policy().action_space,
-                    "initial_epsilon": 0.0,  # <- no randomness whatsoever
-                }
-            }
-            config["explore"] = True
-            trainer = dqn.DQNTrainer(config=config, env="FrozenLake-v0")
-            # Now, when we act - even with explore=True - we would expect
-            # the same action for the same input (parameter noise is
-            # deterministic).
-            policy = trainer.get_policy()
-            p_sess = getattr(policy, "_sess", None)
-            policy.exploration.on_episode_start(policy, tf_sess=p_sess)
-            a_ = trainer.compute_action(obs)
-            for _ in range(10):
-                a = trainer.compute_action(obs, explore=True)
-                check(a, a_)
-
-    def _get_current_noise(self, policy, fw):
-        # If noise not even created yet, return 0.0.
-        if policy.exploration.noise is None:
-            return 0.0
-
-        noise = policy.exploration.noise[0][0][0]
-        if fw == "tf":
-            noise = policy.get_session().run(noise)
-        else:
-            noise = noise.numpy()
-        return noise
-
-    def _get_current_weight(self, policy, fw):
-        weights = policy.get_weights()
-        key = 0 if fw == "eager" else list(weights.keys())[0]
-        return weights[key][0][0]
-
 
 if __name__ == "__main__":
     import pytest

rllib/agents/sac/sac_policy.py

@@ -5,19 +5,18 @@ import ray
 import ray.experimental.tf_utils
 from gym.spaces import Box, Discrete
 from ray.rllib.agents.ddpg.noop_model import NoopModel
-from ray.rllib.agents.dqn.dqn_tf_policy import postprocess_nstep_and_prio, \
-    PRIO_WEIGHTS
+from ray.rllib.agents.ddpg.ddpg_policy import ComputeTDErrorMixin, \
+    TargetNetworkMixin
+from ray.rllib.agents.dqn.dqn_tf_policy import postprocess_nstep_and_prio
 from ray.rllib.agents.sac.sac_model import SACModel
 from ray.rllib.models import ModelCatalog
 from ray.rllib.models.tf.tf_action_dist import (Categorical, SquashedGaussian,
                                                 DiagGaussian)
 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.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
+from ray.rllib.utils.tf_ops import minimize_and_clip
 
 tf = try_import_tf()
 tfp = try_import_tfp()
@@ -117,7 +116,7 @@ def get_distribution_inputs_and_class(policy,
     return distribution_inputs, action_dist_class, state_out
 
 
-def actor_critic_loss(policy, model, _, train_batch):
+def sac_actor_critic_loss(policy, model, _, train_batch):
     model_out_t, _ = model({
         "obs": train_batch[SampleBatch.CUR_OBS],
         "is_training": policy._get_is_training_placeholder(),
@@ -283,7 +282,7 @@ def actor_critic_loss(policy, model, _, train_batch):
 def gradients(policy, optimizer, loss):
     if policy.config["grad_norm_clipping"]:
         actor_grads_and_vars = minimize_and_clip(
-            optimizer,
+            optimizer,  # isn't optimizer not well defined here (which one)?
             policy.actor_loss,
             var_list=policy.model.policy_variables(),
             clip_val=policy.config["grad_norm_clipping"])
@@ -399,63 +398,12 @@ class ActorCriticOptimizerMixin:
             learning_rate=config["optimization"]["entropy_learning_rate"])
 
 
-class ComputeTDErrorMixin:
-    def __init__(self):
-        @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 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),
-                    SampleBatch.ACTIONS: tf.convert_to_tensor(act_t),
-                    SampleBatch.REWARDS: tf.convert_to_tensor(rew_t),
-                    SampleBatch.NEXT_OBS: tf.convert_to_tensor(obs_tp1),
-                    SampleBatch.DONES: tf.convert_to_tensor(done_mask),
-                    PRIO_WEIGHTS: tf.convert_to_tensor(importance_weights),
-                })
-
-            return self.td_error
-
-        self.compute_td_error = compute_td_error
-
-
-class TargetNetworkMixin:
-    def __init__(self, config):
-        @make_tf_callable(self.get_session())
-        def update_target_fn(tau):
-            tau = tf.convert_to_tensor(tau, dtype=tf.float32)
-            update_target_expr = []
-            model_vars = self.model.trainable_variables()
-            target_model_vars = self.target_model.trainable_variables()
-            assert len(model_vars) == len(target_model_vars), \
-                (model_vars, target_model_vars)
-            for var, var_target in zip(model_vars, target_model_vars):
-                update_target_expr.append(
-                    var_target.assign(tau * var + (1.0 - tau) * var_target))
-                logger.debug("Update target op {}".format(var_target))
-            return tf.group(*update_target_expr)
-
-        # Hard initial update
-        self._do_update = update_target_fn
-        self.update_target(tau=1.0)
-
-    # support both hard and soft sync
-    def update_target(self, tau=None):
-        self._do_update(np.float32(tau or self.config.get("tau")))
-
-    @override(TFPolicy)
-    def variables(self):
-        return self.model.variables() + self.target_model.variables()
-
-
 def setup_early_mixins(policy, obs_space, action_space, config):
     ActorCriticOptimizerMixin.__init__(policy, config)
 
 
 def setup_mid_mixins(policy, obs_space, action_space, config):
-    ComputeTDErrorMixin.__init__(policy)
+    ComputeTDErrorMixin.__init__(policy, sac_actor_critic_loss)
 
 
 def setup_late_mixins(policy, obs_space, action_space, config):
@@ -468,7 +416,7 @@ SACTFPolicy = build_tf_policy(
     make_model=build_sac_model,
     postprocess_fn=postprocess_trajectory,
     action_distribution_fn=get_distribution_inputs_and_class,
-    loss_fn=actor_critic_loss,
+    loss_fn=sac_actor_critic_loss,
     stats_fn=stats,
     gradients_fn=gradients,
     apply_gradients_fn=apply_gradients,

rllib/evaluation/rollout_worker.py

@@ -254,11 +254,11 @@ class RolloutWorker(EvaluatorInterface, ParallelIteratorWorker):
 
         policy_config = policy_config or {}
         if (tf and policy_config.get("eager")
-                and not policy_config.get("no_eager_on_workers")):
-            # This check is necessary for certain all-framework tests that
-            # use tf's eager_mode() context generator.
-            if not tf.executing_eagerly():
-                tf.enable_eager_execution()
+                and not policy_config.get("no_eager_on_workers")
+                # This eager check is necessary for certain all-framework tests
+                # that use tf's eager_mode() context generator.
+                and not tf.executing_eagerly()):
+            tf.enable_eager_execution()
 
         if log_level:
             logging.getLogger("ray.rllib").setLevel(log_level)

rllib/models/modelv2.py

@@ -106,11 +106,13 @@ class ModelV2:
         You can find an runnable example in examples/custom_loss.py.
 
         Arguments:
-            policy_loss (Tensor): scalar policy loss from the policy.
+            policy_loss (Union[List[Tensor],Tensor]): List of or single policy
+                loss(es) from the policy.
             loss_inputs (dict): map of input placeholders for rollout data.
 
         Returns:
-            Scalar tensor for the customized loss for this model.
+            Union[List[Tensor],Tensor]: List of or scalar tensor for the
+                customized loss(es) for this model.
         """
         return policy_loss
 

rllib/models/torch/torch_action_dist.py

@@ -157,3 +157,28 @@ class TorchDiagGaussian(TorchDistributionWrapper):
     @override(ActionDistribution)
     def required_model_output_shape(action_space, model_config):
         return np.prod(action_space.shape) * 2
+
+
+class TorchDeterministic(TorchDistributionWrapper):
+    """Action distribution that returns the input values directly.
+
+    This is similar to DiagGaussian with standard deviation zero (thus only
+    requiring the "mean" values as NN output).
+    """
+
+    @override(ActionDistribution)
+    def deterministic_sample(self):
+        return self.inputs
+
+    @override(TorchDistributionWrapper)
+    def sampled_action_logp(self):
+        return 0.0
+
+    @override(TorchDistributionWrapper)
+    def sample(self):
+        return self.deterministic_sample()
+
+    @staticmethod
+    @override(ActionDistribution)
+    def required_model_output_shape(action_space, model_config):
+        return np.prod(action_space.shape)

rllib/train.py

@@ -156,10 +156,10 @@ def run(args, parser):
 
     verbose = 1
     for exp in experiments.values():
-
         # Bazel makes it hard to find files specified in `args` (and `data`).
         # Look for them here.
-        if exp["config"].get("input") and \
+        # NOTE: Some of our yaml files don't have a `config` section.
+        if exp.get("config", {}).get("input") and \
                 not os.path.exists(exp["config"]["input"]):
             # This script runs in the ray/rllib dir.
             rllib_dir = Path(__file__).parent

rllib/utils/exploration/gaussian_noise.py

@@ -163,4 +163,5 @@ class GaussianNoise(Exploration):
         Returns:
             Union[float,tf.Tensor[float]]: The current scale value.
         """
-        return self.scale_schedule(self.last_timestep)
+        scale = self.scale_schedule(self.last_timestep)
+        return {"cur_scale": scale}

rllib/utils/exploration/parameter_noise.py

@@ -1,10 +1,11 @@
-from gym.spaces import Discrete
+from gym.spaces import Box, Discrete
 import numpy as np
 
 from ray.rllib.policy.sample_batch import SampleBatch
 from ray.rllib.models.modelv2 import ModelV2
-from ray.rllib.models.tf.tf_action_dist import Categorical
-from ray.rllib.models.torch.torch_action_dist import TorchCategorical
+from ray.rllib.models.tf.tf_action_dist import Categorical, Deterministic
+from ray.rllib.models.torch.torch_action_dist import TorchCategorical, \
+    TorchDeterministic
 from ray.rllib.utils.annotations import override
 from ray.rllib.utils.exploration.exploration import Exploration
 from ray.rllib.utils.framework import try_import_tf, try_import_torch
@@ -112,6 +113,11 @@ class ParameterNoise(Exploration):
                         "outside_value": 0.01
                     }
                 }
+            elif isinstance(self.action_space, Box):
+                sub_exploration = {
+                    "type": "OrnsteinUhlenbeckNoise",
+                    "random_timesteps": random_timesteps,
+                }
             # TODO(sven): Implement for any action space.
             else:
                 raise NotImplementedError
@@ -201,6 +207,8 @@ class ParameterNoise(Exploration):
         noisy_action_dist = noise_free_action_dist = None
         # Adjust the stddev depending on the action (pi)-distance.
         # Also see [1] for details.
+        # TODO(sven): Find out whether this can be scrapped by simply using
+        #  the `sample_batch` to get the noisy/noise-free action dist.
         _, _, fetches = policy.compute_actions(
             obs_batch=sample_batch[SampleBatch.CUR_OBS],
             # TODO(sven): What about state-ins and seq-lens?
@@ -211,8 +219,11 @@ class ParameterNoise(Exploration):
         # Categorical case (e.g. DQN).
         if policy.dist_class in (Categorical, TorchCategorical):
             action_dist = softmax(fetches[SampleBatch.ACTION_DIST_INPUTS])
-        else:  # TODO(sven): Other action-dist cases.
-            raise NotImplementedError
+        # Deterministic (Gaussian actions, e.g. DDPG).
+        elif policy.dist_class in [Deterministic, TorchDeterministic]:
+            action_dist = fetches[SampleBatch.ACTION_DIST_INPUTS]
+        else:
+            raise NotImplementedError  # TODO(sven): Other action-dist cases.
 
         if self.weights_are_currently_noisy:
             noisy_action_dist = action_dist
@@ -221,7 +232,6 @@ class ParameterNoise(Exploration):
 
         _, _, fetches = policy.compute_actions(
             obs_batch=sample_batch[SampleBatch.CUR_OBS],
-            # TODO(sven): What about state-ins and seq-lens?
             prev_action_batch=sample_batch.get(SampleBatch.PREV_ACTIONS),
             prev_reward_batch=sample_batch.get(SampleBatch.PREV_REWARDS),
             explore=not self.weights_are_currently_noisy)
@@ -229,18 +239,22 @@ class ParameterNoise(Exploration):
         # Categorical case (e.g. DQN).
         if policy.dist_class in (Categorical, TorchCategorical):
             action_dist = softmax(fetches[SampleBatch.ACTION_DIST_INPUTS])
+            # Deterministic (Gaussian actions, e.g. DDPG).
+        elif policy.dist_class in [Deterministic, TorchDeterministic]:
+            action_dist = fetches[SampleBatch.ACTION_DIST_INPUTS]
 
         if noisy_action_dist is None:
             noisy_action_dist = action_dist
         else:
             noise_free_action_dist = action_dist
 
+        delta = distance = None
         # Categorical case (e.g. DQN).
         if policy.dist_class in (Categorical, TorchCategorical):
             # Calculate KL-divergence (DKL(clean||noisy)) according to [2].
             # TODO(sven): Allow KL-divergence to be calculated by our
             #  Distribution classes (don't support off-graph/numpy yet).
-            kl_divergence = np.nanmean(
+            distance = np.nanmean(
                 np.sum(
                     noise_free_action_dist *
                     np.log(noise_free_action_dist /
@@ -250,10 +264,21 @@ class ParameterNoise(Exploration):
                 current_epsilon = tf_sess.run(current_epsilon)
             delta = -np.log(1 - current_epsilon +
                             current_epsilon / self.action_space.n)
-            if kl_divergence <= delta:
-                self.stddev_val *= 1.01
-            else:
-                self.stddev_val /= 1.01
+        elif policy.dist_class in [Deterministic, TorchDeterministic]:
+            # Calculate MSE between noisy and non-noisy output (see [2]).
+            distance = np.sqrt(
+                np.mean(np.square(noise_free_action_dist - noisy_action_dist)))
+            current_scale = self.sub_exploration.get_info()["cur_scale"]
+            if tf_sess is not None:
+                current_scale = tf_sess.run(current_scale)
+            delta = getattr(self.sub_exploration, "ou_sigma", 0.2) * \
+                current_scale
+
+        # Adjust stddev according to the calculated action-distance.
+        if distance <= delta:
+            self.stddev_val *= 1.01
+        else:
+            self.stddev_val /= 1.01
 
         # Set self.stddev to calculated value.
         if self.framework == "tf":

rllib/utils/exploration/tests/test_parameter_noise.py

@@ -0,0 +1,203 @@
+import numpy as np
+import unittest
+
+import ray.rllib.agents.ddpg as ddpg
+import ray.rllib.agents.dqn as dqn
+from ray.rllib.utils.framework import try_import_tf
+from ray.rllib.utils.test_utils import check, framework_iterator
+
+tf = try_import_tf()
+
+
+class TestParameterNoise(unittest.TestCase):
+    def test_ddpg_parameter_noise(self):
+        self.do_test_parameter_noise_exploration(
+            ddpg.DDPGTrainer,
+            ddpg.DEFAULT_CONFIG,
+            "Pendulum-v0", {},
+            np.array([1.0, 0.0, -1.0]),
+            fws="tf")
+
+    def test_dqn_parameter_noise(self):
+        self.do_test_parameter_noise_exploration(
+            dqn.DQNTrainer,
+            dqn.DEFAULT_CONFIG,
+            "FrozenLake-v0", {
+                "is_slippery": False,
+                "map_name": "4x4"
+            },
+            np.array(0),
+            fws=("tf", "eager"))
+
+    def do_test_parameter_noise_exploration(self, trainer_cls, config, env,
+                                            env_config, obs, fws):
+        """Tests, whether an Agent works with ParameterNoise."""
+        core_config = config.copy()
+        core_config["num_workers"] = 0  # Run locally.
+        core_config["env_config"] = env_config
+
+        for fw in framework_iterator(core_config, fws):
+
+            config = core_config.copy()
+
+            # DQN with ParameterNoise exploration (config["explore"]=True).
+            # ----
+            config["exploration_config"] = {"type": "ParameterNoise"}
+            config["explore"] = True
+
+            trainer = trainer_cls(config=config, env=env)
+            policy = trainer.get_policy()
+            self.assertFalse(policy.exploration.weights_are_currently_noisy)
+            noise_before = self._get_current_noise(policy, fw)
+            check(noise_before, 0.0)
+            initial_weights = self._get_current_weight(policy, fw)
+
+            # Pseudo-start an episode and compare the weights before and after.
+            policy.exploration.on_episode_start(policy, tf_sess=policy._sess)
+            self.assertFalse(policy.exploration.weights_are_currently_noisy)
+            noise_after_ep_start = self._get_current_noise(policy, fw)
+            weights_after_ep_start = self._get_current_weight(policy, fw)
+            # Should be the same, as we don't do anything at the beginning of
+            # the episode, only one step later.
+            check(noise_after_ep_start, noise_before)
+            check(initial_weights, weights_after_ep_start)
+
+            # Setting explore=False should always return the same action.
+            a_ = trainer.compute_action(obs, explore=False)
+            self.assertFalse(policy.exploration.weights_are_currently_noisy)
+            noise = self._get_current_noise(policy, fw)
+            # We sampled the first noise (not zero anymore).
+            check(noise, 0.0, false=True)
+            # But still not applied b/c explore=False.
+            check(self._get_current_weight(policy, fw), initial_weights)
+            for _ in range(10):
+                a = trainer.compute_action(obs, explore=False)
+                check(a, a_)
+                # Noise never gets applied.
+                check(self._get_current_weight(policy, fw), initial_weights)
+                self.assertFalse(
+                    policy.exploration.weights_are_currently_noisy)
+
+            # Explore=None (default: True) should return different actions.
+            # However, this is only due to the underlying epsilon-greedy
+            # exploration.
+            actions = []
+            current_weight = None
+            for _ in range(10):
+                actions.append(trainer.compute_action(obs))
+                self.assertTrue(policy.exploration.weights_are_currently_noisy)
+                # Now, noise actually got applied (explore=True).
+                current_weight = self._get_current_weight(policy, fw)
+                check(current_weight, initial_weights, false=True)
+                check(current_weight, initial_weights + noise)
+            check(np.std(actions), 0.0, false=True)
+
+            # Pseudo-end the episode and compare weights again.
+            # Make sure they are the original ones.
+            policy.exploration.on_episode_end(policy, tf_sess=policy._sess)
+            weights_after_ep_end = self._get_current_weight(policy, fw)
+            check(current_weight - noise, weights_after_ep_end, decimals=5)
+
+            # DQN with ParameterNoise exploration (config["explore"]=False).
+            # ----
+            config = core_config.copy()
+            config["exploration_config"] = {"type": "ParameterNoise"}
+            config["explore"] = False
+            trainer = trainer_cls(config=config, env=env)
+            policy = trainer.get_policy()
+            self.assertFalse(policy.exploration.weights_are_currently_noisy)
+            initial_weights = self._get_current_weight(policy, fw)
+
+            # Noise before anything (should be 0.0, no episode started yet).
+            noise = self._get_current_noise(policy, fw)
+            check(noise, 0.0)
+
+            # Pseudo-start an episode and compare the weights before and after
+            # (they should be the same).
+            policy.exploration.on_episode_start(policy, tf_sess=policy._sess)
+            self.assertFalse(policy.exploration.weights_are_currently_noisy)
+
+            # Should be the same, as we don't do anything at the beginning of
+            # the episode, only one step later.
+            noise = self._get_current_noise(policy, fw)
+            check(noise, 0.0)
+            noisy_weights = self._get_current_weight(policy, fw)
+            check(initial_weights, noisy_weights)
+
+            # Setting explore=False or None should always return the same
+            # action.
+            a_ = trainer.compute_action(obs, explore=False)
+            # Now we have re-sampled.
+            noise = self._get_current_noise(policy, fw)
+            check(noise, 0.0, false=True)
+            for _ in range(5):
+                a = trainer.compute_action(obs, explore=None)
+                check(a, a_)
+                a = trainer.compute_action(obs, explore=False)
+                check(a, a_)
+
+            # Pseudo-end the episode and compare weights again.
+            # Make sure they are the original ones (no noise permanently
+            # applied throughout the episode).
+            policy.exploration.on_episode_end(policy, tf_sess=policy._sess)
+            weights_after_episode_end = self._get_current_weight(policy, fw)
+            check(initial_weights, weights_after_episode_end)
+            # Noise should still be the same (re-sampling only happens at
+            # beginning of episode).
+            noise_after = self._get_current_noise(policy, fw)
+            check(noise, noise_after)
+
+            # Switch off underlying exploration entirely.
+            # ----
+            config = core_config.copy()
+            if trainer_cls is dqn.DQNTrainer:
+                sub_config = {
+                    "type": "EpsilonGreedy",
+                    "initial_epsilon": 0.0,  # <- no randomness whatsoever
+                    "final_epsilon": 0.0,
+                }
+            else:
+                sub_config = {
+                    "type": "OrnsteinUhlenbeckNoise",
+                    "initial_scale": 0.0,  # <- no randomness whatsoever
+                    "final_scale": 0.0,
+                    "random_timesteps": 0,
+                }
+            config["exploration_config"] = {
+                "type": "ParameterNoise",
+                "sub_exploration": sub_config,
+            }
+            config["explore"] = True
+            trainer = trainer_cls(config=config, env=env)
+            # Now, when we act - even with explore=True - we would expect
+            # the same action for the same input (parameter noise is
+            # deterministic).
+            policy = trainer.get_policy()
+            policy.exploration.on_episode_start(policy, tf_sess=policy._sess)
+            a_ = trainer.compute_action(obs)
+            for _ in range(10):
+                a = trainer.compute_action(obs, explore=True)
+                check(a, a_)
+
+    def _get_current_noise(self, policy, fw):
+        # If noise not even created yet, return 0.0.
+        if policy.exploration.noise is None:
+            return 0.0
+
+        noise = policy.exploration.noise[0][0][0]
+        if fw == "tf":
+            noise = policy.get_session().run(noise)
+        else:
+            noise = noise.numpy()
+        return noise
+
+    def _get_current_weight(self, policy, fw):
+        weights = policy.get_weights()
+        key = 0 if fw == "eager" else list(weights.keys())[0]
+        return weights[key][0][0]
+
+
+if __name__ == "__main__":
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))