Allow usage with tensorflow 2.0.0 (via tf.compat.v1) (#2665)

Chris Elion · web-flow · commit 182a40e6f7ce · 2019-11-13T18:23:57.000-08:00
diff --git a/.circleci/config.yml b/.circleci/config.yml
@@ -228,7 +228,13 @@ workflows:
           executor: python373
           pyversion: 3.7.3
           # Test python 3.7 with the newest supported versions
-          pip_constraints: test_constraints_max_version.txt
+          pip_constraints: test_constraints_max_tf1_version.txt
+      - build_python:
+          name: python_3.7.3+tf2
+          executor: python373
+          pyversion: 3.7.3
+          # Test python 3.7 with the newest supported versions
+          pip_constraints: test_constraints_max_tf2_version.txt
       - markdown_link_check
       - protobuf_generation_check
       - deploy:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -44,7 +44,8 @@ repos:
                 .*_pb2.py|
                 .*_pb2_grpc.py
             )$
-        additional_dependencies: [flake8-comprehensions]
+        # flake8-tidy-imports is used for banned-modules, not actually tidying
+        additional_dependencies: [flake8-comprehensions, flake8-tidy-imports]
     -   id: trailing-whitespace
         name: trailing-whitespace-markdown
         types: [markdown]
diff --git a/ml-agents/mlagents/tf_utils/__init__.py b/ml-agents/mlagents/tf_utils/__init__.py
@@ -0,0 +1 @@
+from mlagents.tf_utils.tf import tf as tf  # noqa
diff --git a/ml-agents/mlagents/tf_utils/tf.py b/ml-agents/mlagents/tf_utils/tf.py
@@ -0,0 +1,15 @@
+# This should be the only place that we import tensorflow directly.
+# Everywhere else is caught by the banned-modules setting for flake8
+import tensorflow as tf  # noqa I201
+from distutils.version import LooseVersion
+
+
+# LooseVersion handles things "1.2.3a" or "4.5.6-rc7" fairly sensibly.
+_is_tensorflow2 = LooseVersion(tf.__version__) >= LooseVersion("2.0.0")
+
+if _is_tensorflow2:
+    import tensorflow.compat.v1 as tf
+
+    tf.disable_v2_behavior()
+else:
+    pass
diff --git a/ml-agents/mlagents/trainers/bc/models.py b/ml-agents/mlagents/trainers/bc/models.py
@@ -1,5 +1,5 @@
-import tensorflow as tf
-import tensorflow.contrib.layers as c_layers
+from mlagents.tf_utils import tf
+
 from mlagents.trainers.models import LearningModel
 
 
@@ -44,9 +44,7 @@ def __init__(
                         size,
                         activation=None,
                         use_bias=False,
-                        kernel_initializer=c_layers.variance_scaling_initializer(
-                            factor=0.01
-                        ),
+                        kernel_initializer=tf.initializers.variance_scaling(0.01),
                     )
                 )
             self.action_probs = tf.concat(
@@ -93,7 +91,7 @@ def __init__(
                 activation=None,
                 use_bias=False,
                 name="pre_action",
-                kernel_initializer=c_layers.variance_scaling_initializer(factor=0.01),
+                kernel_initializer=tf.initializers.variance_scaling(0.01),
             )
             self.clipped_sample_action = tf.clip_by_value(self.policy, -1, 1)
             self.sample_action = tf.identity(self.clipped_sample_action, name="action")
diff --git a/ml-agents/mlagents/trainers/components/bc/model.py b/ml-agents/mlagents/trainers/components/bc/model.py
@@ -1,4 +1,5 @@
-import tensorflow as tf
+from mlagents.tf_utils import tf
+
 from mlagents.trainers.models import LearningModel
 
 
diff --git a/ml-agents/mlagents/trainers/components/reward_signals/__init__.py b/ml-agents/mlagents/trainers/components/reward_signals/__init__.py
@@ -4,7 +4,7 @@
 import numpy as np
 import abc
 
-import tensorflow as tf
+from mlagents.tf_utils import tf
 
 from mlagents.envs.brain import BrainInfo
 from mlagents.trainers.trainer import UnityTrainerException
diff --git a/ml-agents/mlagents/trainers/components/reward_signals/curiosity/model.py b/ml-agents/mlagents/trainers/components/reward_signals/curiosity/model.py
@@ -1,5 +1,6 @@
 from typing import List, Tuple
-import tensorflow as tf
+from mlagents.tf_utils import tf
+
 from mlagents.trainers.models import LearningModel
 
 
diff --git a/ml-agents/mlagents/trainers/components/reward_signals/curiosity/signal.py b/ml-agents/mlagents/trainers/components/reward_signals/curiosity/signal.py
@@ -1,6 +1,7 @@
 from typing import Any, Dict, List
 import numpy as np
-import tensorflow as tf
+from mlagents.tf_utils import tf
+
 from mlagents.envs.brain import BrainInfo
 
 from mlagents.trainers.components.reward_signals import RewardSignal, RewardSignalResult
diff --git a/ml-agents/mlagents/trainers/components/reward_signals/gail/model.py b/ml-agents/mlagents/trainers/components/reward_signals/gail/model.py
@@ -1,6 +1,7 @@
 from typing import List, Optional, Tuple
 
-import tensorflow as tf
+from mlagents.tf_utils import tf
+
 from mlagents.trainers.models import LearningModel
 
 EPSILON = 1e-7
diff --git a/ml-agents/mlagents/trainers/components/reward_signals/gail/signal.py b/ml-agents/mlagents/trainers/components/reward_signals/gail/signal.py
@@ -1,7 +1,7 @@
 from typing import Any, Dict, List
 import logging
 import numpy as np
-import tensorflow as tf
+from mlagents.tf_utils import tf
 
 from mlagents.envs.brain import BrainInfo
 from mlagents.trainers.components.reward_signals import RewardSignal, RewardSignalResult
diff --git a/ml-agents/mlagents/trainers/models.py b/ml-agents/mlagents/trainers/models.py
@@ -3,8 +3,7 @@
 from typing import Callable, Dict, List, Optional
 
 import numpy as np
-import tensorflow as tf
-import tensorflow.contrib.layers as c_layers
+from mlagents.tf_utils import tf
 
 from mlagents.trainers.trainer import UnityTrainerException
 from mlagents.envs.brain import CameraResolution
@@ -123,7 +122,7 @@ def create_learning_rate(
 
     @staticmethod
     def scaled_init(scale):
-        return c_layers.variance_scaling_initializer(scale)
+        return tf.initializers.variance_scaling(scale)
 
     @staticmethod
     def swish(input_activation: tf.Tensor) -> tf.Tensor:
@@ -245,7 +244,7 @@ def create_vector_observation_encoder(
                     activation=activation,
                     reuse=reuse,
                     name="hidden_{}".format(i),
-                    kernel_initializer=c_layers.variance_scaling_initializer(1.0),
+                    kernel_initializer=tf.initializers.variance_scaling(1.0),
                 )
         return hidden
 
@@ -287,7 +286,7 @@ def create_visual_observation_encoder(
                 reuse=reuse,
                 name="conv_2",
             )
-            hidden = c_layers.flatten(conv2)
+            hidden = tf.layers.flatten(conv2)
 
         with tf.variable_scope(scope + "/" + "flat_encoding"):
             hidden_flat = LearningModel.create_vector_observation_encoder(
@@ -342,7 +341,7 @@ def create_nature_cnn_visual_observation_encoder(
                 reuse=reuse,
                 name="conv_3",
             )
-            hidden = c_layers.flatten(conv3)
+            hidden = tf.layers.flatten(conv3)
 
         with tf.variable_scope(scope + "/" + "flat_encoding"):
             hidden_flat = LearningModel.create_vector_observation_encoder(
@@ -410,7 +409,7 @@ def create_resnet_visual_observation_encoder(
                     )
                     hidden = tf.add(block_input, hidden)
             hidden = tf.nn.relu(hidden)
-            hidden = c_layers.flatten(hidden)
+            hidden = tf.layers.flatten(hidden)
 
         with tf.variable_scope(scope + "/" + "flat_encoding"):
             hidden_flat = LearningModel.create_vector_observation_encoder(
@@ -557,8 +556,8 @@ def create_recurrent_encoder(input_state, memory_in, sequence_length, name="lstm
         memory_in = tf.reshape(memory_in[:, :], [-1, m_size])
         half_point = int(m_size / 2)
         with tf.variable_scope(name):
-            rnn_cell = tf.contrib.rnn.BasicLSTMCell(half_point)
-            lstm_vector_in = tf.contrib.rnn.LSTMStateTuple(
+            rnn_cell = tf.nn.rnn_cell.BasicLSTMCell(half_point)
+            lstm_vector_in = tf.nn.rnn_cell.LSTMStateTuple(
                 memory_in[:, :half_point], memory_in[:, half_point:]
             )
             recurrent_output, lstm_state_out = tf.nn.dynamic_rnn(
diff --git a/ml-agents/mlagents/trainers/ppo/models.py b/ml-agents/mlagents/trainers/ppo/models.py
@@ -2,8 +2,7 @@
 from typing import Optional
 
 import numpy as np
-import tensorflow as tf
-
+from mlagents.tf_utils import tf
 from mlagents.trainers.models import LearningModel, EncoderType, LearningRateSchedule
 
 logger = logging.getLogger("mlagents.trainers")
diff --git a/ml-agents/mlagents/trainers/ppo/multi_gpu_policy.py b/ml-agents/mlagents/trainers/ppo/multi_gpu_policy.py
@@ -1,7 +1,8 @@
 import logging
 from typing import Any, Dict, List, Optional
 
-import tensorflow as tf
+from mlagents.tf_utils import tf
+
 from tensorflow.python.client import device_lib
 from mlagents.envs.brain import BrainParameters
 from mlagents.envs.timers import timed
diff --git a/ml-agents/mlagents/trainers/ppo/policy.py b/ml-agents/mlagents/trainers/ppo/policy.py
@@ -1,7 +1,8 @@
 import logging
 import numpy as np
 from typing import Any, Dict, Optional
-import tensorflow as tf
+
+from mlagents.tf_utils import tf
 
 from mlagents.envs.timers import timed
 from mlagents.envs.brain import BrainInfo, BrainParameters
diff --git a/ml-agents/mlagents/trainers/sac/models.py b/ml-agents/mlagents/trainers/sac/models.py
@@ -2,9 +2,9 @@
 import numpy as np
 from typing import Dict, List, Optional
 
-import tensorflow as tf
+from mlagents.tf_utils import tf
+
 from mlagents.trainers.models import LearningModel, LearningRateSchedule, EncoderType
-import tensorflow.contrib.layers as c_layers
 
 LOG_STD_MAX = 2
 LOG_STD_MIN = -20
@@ -313,9 +313,7 @@ def create_dc_actor(self, hidden_policy, scope):
                         size,
                         activation=None,
                         use_bias=False,
-                        kernel_initializer=c_layers.variance_scaling_initializer(
-                            factor=0.01
-                        ),
+                        kernel_initializer=tf.initializers.variance_scaling(0.01),
                     )
                 )
             all_logits = tf.concat(
diff --git a/ml-agents/mlagents/trainers/sac/policy.py b/ml-agents/mlagents/trainers/sac/policy.py
@@ -1,7 +1,7 @@
 import logging
 from typing import Dict, Any, Optional
 import numpy as np
-import tensorflow as tf
+from mlagents.tf_utils import tf
 
 from mlagents.envs.timers import timed
 from mlagents.envs.brain import BrainInfo, BrainParameters
diff --git a/ml-agents/mlagents/trainers/tensorflow_to_barracuda.py b/ml-agents/mlagents/trainers/tensorflow_to_barracuda.py
@@ -2,7 +2,7 @@
 from __future__ import print_function
 import numpy as np
 import struct  # convert from Python values and C structs
-import tensorflow as tf
+from mlagents.tf_utils import tf
 import re
 
 # import barracuda
diff --git a/ml-agents/mlagents/trainers/tests/test_bc.py b/ml-agents/mlagents/trainers/tests/test_bc.py
@@ -3,7 +3,7 @@
 import os
 
 import numpy as np
-import tensorflow as tf
+from mlagents.tf_utils import tf
 import yaml
 
 from mlagents.trainers.bc.models import BehavioralCloningModel
diff --git a/ml-agents/mlagents/trainers/tests/test_multigpu.py b/ml-agents/mlagents/trainers/tests/test_multigpu.py
@@ -1,7 +1,7 @@
 import unittest.mock as mock
 import pytest
 
-import tensorflow as tf
+from mlagents.tf_utils import tf
 import yaml
 
 from mlagents.trainers.ppo.multi_gpu_policy import MultiGpuPPOPolicy
diff --git a/ml-agents/mlagents/trainers/tests/test_ppo.py b/ml-agents/mlagents/trainers/tests/test_ppo.py
@@ -2,7 +2,8 @@
 import pytest
 
 import numpy as np
-import tensorflow as tf
+from mlagents.tf_utils import tf
+
 import yaml
 
 from mlagents.trainers.ppo.models import PPOModel
diff --git a/ml-agents/mlagents/trainers/tests/test_sac.py b/ml-agents/mlagents/trainers/tests/test_sac.py
@@ -3,7 +3,8 @@
 import yaml
 
 import numpy as np
-import tensorflow as tf
+from mlagents.tf_utils import tf
+
 
 from mlagents.trainers.sac.models import SACModel
 from mlagents.trainers.sac.policy import SACPolicy
diff --git a/ml-agents/mlagents/trainers/tests/test_trainer_controller.py b/ml-agents/mlagents/trainers/tests/test_trainer_controller.py
@@ -1,5 +1,7 @@
 from unittest.mock import MagicMock, Mock, patch
 
+from mlagents.tf_utils import tf
+
 import yaml
 import pytest
 
@@ -56,7 +58,7 @@ def basic_trainer_controller():
 
 
 @patch("numpy.random.seed")
-@patch("tensorflow.set_random_seed")
+@patch.object(tf, "set_random_seed")
 def test_initialization_seed(numpy_random_seed, tensorflow_set_seed):
     seed = 27
     TrainerController(
@@ -102,7 +104,7 @@ def take_step_sideeffect(env):
     return tc, trainer_mock
 
 
-@patch("tensorflow.reset_default_graph")
+@patch.object(tf, "reset_default_graph")
 def test_start_learning_trains_forever_if_no_train_model(tf_reset_graph):
     tc, trainer_mock = trainer_controller_with_start_learning_mocks()
     tc.train_model = False
@@ -123,7 +125,7 @@ def test_start_learning_trains_forever_if_no_train_model(tf_reset_graph):
     env_mock.close.assert_called_once()
 
 
-@patch("tensorflow.reset_default_graph")
+@patch.object(tf, "reset_default_graph")
 def test_start_learning_trains_until_max_steps_then_saves(tf_reset_graph):
     tc, trainer_mock = trainer_controller_with_start_learning_mocks()
     tf_reset_graph.return_value = None
diff --git a/ml-agents/mlagents/trainers/tf_policy.py b/ml-agents/mlagents/trainers/tf_policy.py
@@ -2,7 +2,7 @@
 from typing import Any, Dict, List, Optional
 
 import numpy as np
-import tensorflow as tf
+from mlagents.tf_utils import tf
 
 from mlagents.envs.exception import UnityException
 from mlagents.envs.policy import Policy
diff --git a/ml-agents/mlagents/trainers/trainer.py b/ml-agents/mlagents/trainers/trainer.py
@@ -2,7 +2,9 @@
 import logging
 from typing import Dict, List, Deque, Any
 import os
-import tensorflow as tf
+
+from mlagents.tf_utils import tf
+
 import numpy as np
 from collections import deque, defaultdict
 
diff --git a/ml-agents/mlagents/trainers/trainer_controller.py b/ml-agents/mlagents/trainers/trainer_controller.py
@@ -8,7 +8,7 @@
 from typing import Dict, List, Optional, Set
 
 import numpy as np
-import tensorflow as tf
+from mlagents.tf_utils import tf
 from time import time
 
 from mlagents.envs.env_manager import EnvironmentStep
diff --git a/ml-agents/setup.py b/ml-agents/setup.py
@@ -64,7 +64,7 @@ def run(self):
         "Pillow>=4.2.1",
         "protobuf>=3.6",
         "pyyaml",
-        "tensorflow>=1.7,<2.0",
+        "tensorflow>=1.7,<2.1",
         'pypiwin32==223;platform_system=="Windows"',
     ],
     python_requires=">=3.6.1",
diff --git a/setup.cfg b/setup.cfg
@@ -16,3 +16,7 @@ ignore =
     # Black tends to introduce things flake8 doesn't like, such as "line break before binary operator"
     # or whitespace before ':'. Rather than fight with black, just ignore these for now.
     W503, E203,
+    # flake-tidy-import adds this warning, which we don't really care about for now
+    I200
+
+banned-modules = tensorflow = use mlagents.tf_utils instead (it handles tf2 compat).
diff --git a/test_constraints_max_tf1_version.txt b/test_constraints_max_tf1_version.txt
@@ -0,0 +1,6 @@
+# pip constraints to use the *highest* versions allowed in ml-agents/setup.py
+# with the exception of tensorflow, which is constrained to <2
+# For projects with upper bounds, we should periodically update this list to the latest release version
+grpcio>=1.23.0
+numpy>=1.17.2
+tensorflow>=1.14.0,<2.0
diff --git a/test_constraints_max_tf2_version.txt b/test_constraints_max_tf2_version.txt
@@ -2,4 +2,4 @@
 # For projects with upper bounds, we should periodically update this list to the latest release version
 grpcio>=1.23.0
 numpy>=1.17.2
-tensorflow>=1.14.0,<2.0
+tensorflow>=2.0.0,<2.1.0

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+from mlagents.tf_utils.tf import tf as tf # noqa`