Repository: ray-project/ray

Issue:
[rllib] ModuleNotFoundError: No module named 'tensorflow.contrib'
<!--Please include [tune], [rllib], [autoscaler] etc. in the issue title if relevant-->

### What is the problem?

```
File "/home/axion/anaconda3/envs/training/lib/python3.7/site-packages/ray/rllib/agents/trainer_template.py", line 93, in __init__
    Trainer.__init__(self, config, env, logger_creator)
  File "/home/axion/anaconda3/envs/training/lib/python3.7/site-packages/ray/rllib/agents/trainer.py", line 455, in __init__
    super().__init__(config, logger_creator)
  File "/home/axion/anaconda3/envs/training/lib/python3.7/site-packages/ray/tune/trainable.py", line 174, in __init__
    self._setup(copy.deepcopy(self.config))
  File "/home/axion/anaconda3/envs/training/lib/python3.7/site-packages/ray/rllib/agents/trainer.py", line 596, in _setup
    self._init(self.config, self.env_creator)
  File "/home/axion/anaconda3/envs/training/lib/python3.7/site-packages/ray/rllib/agents/trainer_template.py", line 117, in _init
    config)
  File "/home/axion/anaconda3/envs/training/lib/python3.7/site-packages/ray/rllib/agents/dqn/apex.py", line 48, in defer_make_workers
    return trainer._make_workers(env_creator, policy, config, 0)
  File "/home/axion/anaconda3/envs/training/lib/python3.7/site-packages/ray/rllib/agents/trainer.py", line 666, in _make_workers
    logdir=self.logdir)
  File "/home/axion/anaconda3/envs/training/lib/python3.7/site-packages/ray/rllib/evaluation/worker_set.py", line 61, in __init__
    RolloutWorker, env_creator, policy, 0, self._local_config)
  File "/home/axion/anaconda3/envs/training/lib/python3.7/site-packages/ray/rllib/evaluation/worker_set.py", line 271, in _make_worker
    _fake_sampler=config.get("_fake_sampler", False))
  File "/home/axion/anaconda3/envs/training/lib/python3.7/site-packages/ray/rllib/evaluation/rollout_worker.py", line 360, in __init__
    self._build_policy_map(policy_dict, policy_config)
  File "/home/axion/anaconda3/envs/training/lib/python3.7/site-packages/ray/rllib/evaluation/rollout_worker.py", line 836, in _build_policy_map
    policy_map[name] = cls(obs_space, act_space, merged_conf)
  File "/home/axion/anaconda3/envs/training/lib/python3.7/site-packages/ray/rllib/agents/ddpg/ddpg_policy.py", line 110, in __init__
    self.cur_observations, observation_space, action_space)
  File "/home/axion/anaconda3/envs/training/lib/python3.7/site-packages/ray/rllib/agents/ddpg/ddpg_policy.py", line 410, in _build_policy_network
    import tensorflow.contrib.layers as layers
ModuleNotFoundError: No module named 'tensorflow.contrib'
```

*Ray version and other system information (Python version, TensorFlow version, OS):*
Python 3.7.7
I am running Ray 0.8.3 but this bug is currently still in the master branch
tensorflow==2.1.0

### Reproduction (REQUIRED)
Setting any DDPG algorithm to `parameter_noise=True` forces this import to run
https://github.com/ray-project/ray/blob/master/rllib/agents/ddpg/ddpg_policy.py#L417

**Tensorflow 2 does not have a `contrib` module**

- [x] I have verified my script runs in a clean environment and reproduces the issue.
- [x] I have verified the issue also occurs with the [latest wheels](https://ray.readthedocs.io/en/latest/installation.html).



Discussion:
Yeah, we should switch to keras here. I'm taking a look.
For now, could you replace the code (~ line 416) in rllib/agents/ddpg/ddpg_policy.py with the following? That should eliminate the dependency on the deprecated tf.contrib.

```
        activation = getattr(tf.nn, self.config["actor_hidden_activation"])
        for hidden in self.config["actor_hiddens"]:
            if self.config["parameter_noise"]:
                action_out = tf.keras.layers.Dense(
                    units=hidden,
                    activation=activation)(action_out)
                action_out = tf.keras.layers.LayerNormalization()(action_out)
            else:
                action_out = tf.keras.layers.Dense(
                    units=hidden, activation=activation)(action_out)
        action_out = tf.keras.layers.Dense(
            units=action_space.shape[0], activation=None)(action_out)
```
Cool. I will give this a try!
Preliminary tests look ok. Please let me know, if this doesn't help.

Resolving patch:
diff --git a/rllib/agents/ddpg/ddpg_policy.py b/rllib/agents/ddpg/ddpg_policy.py
--- a/rllib/agents/ddpg/ddpg_policy.py
+++ b/rllib/agents/ddpg/ddpg_policy.py
@@ -54,6 +54,7 @@ def postprocess_trajectory(self,
                 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 @@ def _build_policy_network(self, obs, obs_space, action_space):
 
         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
--- a/rllib/agents/dqn/distributional_q_model.py
+++ b/rllib/agents/dqn/distributional_q_model.py
@@ -125,14 +125,12 @@ def build_state_score(model_out):
                     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",


Test patch: