From 7b08db9f8cd85d185879e5bef778e8855f2a06cf Mon Sep 17 00:00:00 2001
From: Sven Mika <sven@anyscale.io>
Date: Thu, 2 Apr 2020 03:03:14 +0200
Subject: [PATCH] [RLlib] Remove all instances of tf.contrib.layers. ... from
 RLlib code (deprecated). (#7851)

---
 rllib/agents/ddpg/ddpg_policy.py           | 13 ++++---------
 rllib/agents/dqn/distributional_q_model.py |  8 +++-----
 2 files changed, 7 insertions(+), 14 deletions(-)

diff --git a/rllib/agents/ddpg/ddpg_policy.py b/rllib/agents/ddpg/ddpg_policy.py
index f0d54b4a6..641f34284 100644
--- a/rllib/agents/ddpg/ddpg_policy.py
+++ b/rllib/agents/ddpg/ddpg_policy.py
@@ -54,6 +54,7 @@ class DDPGPostprocessing:
                 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)))
@@ -414,16 +415,10 @@ class DDPGTFPolicy(DDPGPostprocessing, TFPolicy):
 
         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"]:
-                import tensorflow.contrib.layers as layers
-                action_out = layers.fully_connected(
-                    action_out,
-                    num_outputs=hidden,
-                    activation_fn=activation,
-                    normalizer_fn=layers.layer_norm)
-            else:
-                action_out = tf.layers.dense(
-                    action_out, units=hidden, activation=activation)
+                action_out = tf.keras.layers.LayerNormalization()(action_out)
         action_out = tf.layers.dense(
             action_out, units=action_space.shape[0], activation=None)
 
diff --git a/rllib/agents/dqn/distributional_q_model.py b/rllib/agents/dqn/distributional_q_model.py
index b331317b0..dd981412b 100644
--- a/rllib/agents/dqn/distributional_q_model.py
+++ b/rllib/agents/dqn/distributional_q_model.py
@@ -125,14 +125,12 @@ class DistributionalQModel(TFModelV2):
                     state_out = self._noisy_layer("dueling_hidden_%d" % i,
                                                   state_out, q_hiddens[i],
                                                   sigma0)
-                elif parameter_noise:
-                    state_out = tf.keras.layers.Dense(
-                        units=q_hiddens[i],
-                        activation_fn=tf.nn.relu,
-                        normalizer_fn=tf.contrib.layers.layer_norm)(state_out)
                 else:
                     state_out = tf.keras.layers.Dense(
                         units=q_hiddens[i], activation=tf.nn.relu)(state_out)
+                    if parameter_noise:
+                        state_out = tf.keras.layers.LayerNormalization()(
+                            state_out)
             if use_noisy:
                 state_score = self._noisy_layer(
                     "dueling_output",