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

.circleci/config.yml

@@ -120,7 +120,7 @@ jobs:
           command: |
             . venv/bin/activate
             pre-commit run --hook-stage manual markdown-link-check --all-files
-  
+
   protobuf_generation_check:
     docker:
       - image: circleci/python:3.7.3
@@ -184,6 +184,12 @@ 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

.pre-commit-config.yaml

@@ -42,6 +42,7 @@ repos:
                 .*_pb2.py|
                 .*_pb2_grpc.py
             )$
+        additional_dependencies: [flake8-tidy-imports]
 # "Local" hooks, see https://pre-commit.com/#repository-local-hooks
 -   repo: local
     hooks:

ml-agents/mlagents/trainers/__init__.py

@@ -0,0 +1,2 @@
+from mlagents.trainers.tf import tf as tf  # noqa
+from mlagents.trainers.tf import tf_flatten, tf_rnn, tf_variance_scaling  # noqa

ml-agents/mlagents/trainers/bc/models.py

@@ -1,5 +1,5 @@
-import tensorflow as tf
-import tensorflow.contrib.layers as c_layers
+from mlagents.trainers import tf, tf_variance_scaling
+
 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_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_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")

ml-agents/mlagents/trainers/components/bc/model.py

@@ -1,4 +1,5 @@
-import tensorflow as tf
+from mlagents.trainers import tf
+
 from mlagents.trainers.models import LearningModel
 
 

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.trainers import tf
 
 from mlagents.envs.brain import BrainInfo
 from mlagents.trainers.trainer import UnityTrainerException

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.trainers import tf
+
 from mlagents.trainers.models import LearningModel
 
 

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.trainers import tf
+
 from mlagents.envs.brain import BrainInfo
 
 from mlagents.trainers.components.reward_signals import RewardSignal, RewardSignalResult

ml-agents/mlagents/trainers/components/reward_signals/gail/model.py

@@ -1,6 +1,7 @@
 from typing import Tuple, List
 
-import tensorflow as tf
+from mlagents.trainers import tf
+
 from mlagents.trainers.models import LearningModel
 
 EPSILON = 1e-7

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.trainers import tf
 
 from mlagents.envs.brain import BrainInfo
 from mlagents.trainers.components.reward_signals import RewardSignal, RewardSignalResult

ml-agents/mlagents/trainers/models.py

@@ -3,8 +3,7 @@
 from typing import Callable, List
 
 import numpy as np
-import tensorflow as tf
-import tensorflow.contrib.layers as c_layers
+from mlagents.trainers import tf, tf_variance_scaling, tf_rnn, tf_flatten
 
 from mlagents.trainers.trainer import UnityTrainerException
 from mlagents.envs.brain import CameraResolution
@@ -119,7 +118,7 @@ def create_learning_rate(
 
     @staticmethod
     def scaled_init(scale):
-        return c_layers.variance_scaling_initializer(scale)
+        return tf_variance_scaling(scale)
 
     @staticmethod
     def swish(input_activation: tf.Tensor) -> tf.Tensor:
@@ -241,7 +240,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_variance_scaling(1.0),
                 )
         return hidden
 
@@ -283,7 +282,7 @@ def create_visual_observation_encoder(
                 reuse=reuse,
                 name="conv_2",
             )
-            hidden = c_layers.flatten(conv2)
+            hidden = tf_flatten(conv2)
 
         with tf.variable_scope(scope + "/" + "flat_encoding"):
             hidden_flat = LearningModel.create_vector_observation_encoder(
@@ -338,7 +337,7 @@ def create_nature_cnn_visual_observation_encoder(
                 reuse=reuse,
                 name="conv_3",
             )
-            hidden = c_layers.flatten(conv3)
+            hidden = tf_flatten(conv3)
 
         with tf.variable_scope(scope + "/" + "flat_encoding"):
             hidden_flat = LearningModel.create_vector_observation_encoder(
@@ -406,7 +405,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_flatten(hidden)
 
         with tf.variable_scope(scope + "/" + "flat_encoding"):
             hidden_flat = LearningModel.create_vector_observation_encoder(
@@ -553,8 +552,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_rnn.BasicLSTMCell(half_point)
+            lstm_vector_in = tf_rnn.LSTMStateTuple(
                 memory_in[:, :half_point], memory_in[:, half_point:]
             )
             recurrent_output, lstm_state_out = tf.nn.dynamic_rnn(

ml-agents/mlagents/trainers/ppo/models.py

@@ -1,7 +1,8 @@
 import logging
 import numpy as np
 
-import tensorflow as tf
+from mlagents.trainers import tf
+
 from mlagents.trainers.models import LearningModel, EncoderType, LearningRateSchedule
 
 logger = logging.getLogger("mlagents.trainers")

ml-agents/mlagents/trainers/ppo/multi_gpu_policy.py

@@ -1,6 +1,7 @@
 import logging
 
-import tensorflow as tf
+from mlagents.trainers import tf
+
 from tensorflow.python.client import device_lib
 from mlagents.envs.timers import timed
 from mlagents.trainers.models import EncoderType, LearningRateSchedule

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.trainers import tf
 
 from mlagents.envs.timers import timed
 from mlagents.envs.brain import BrainInfo, BrainParameters

ml-agents/mlagents/trainers/sac/models.py

@@ -1,9 +1,9 @@
 import logging
 import numpy as np
 
-import tensorflow as tf
+from mlagents.trainers import tf, tf_variance_scaling
+
 from mlagents.trainers.models import LearningModel, LearningRateSchedule, EncoderType
-import tensorflow.contrib.layers as c_layers
 
 LOG_STD_MAX = 2
 LOG_STD_MIN = -20
@@ -275,9 +275,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_variance_scaling(0.01),
                     )
                 )
             all_logits = tf.concat(

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.trainers import tf
 
 from mlagents.envs.timers import timed
 from mlagents.envs.brain import BrainInfo, BrainParameters

ml-agents/mlagents/trainers/tensorflow_to_barracuda.py

@@ -1,7 +1,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.trainers import tf
 import re
 
 # import barracuda

ml-agents/mlagents/trainers/tests/test_bc.py

@@ -3,7 +3,7 @@
 import os
 
 import numpy as np
-import tensorflow as tf
+from mlagents.trainers import tf
 import yaml
 
 from mlagents.trainers.bc.models import BehavioralCloningModel

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.trainers import tf
 import yaml
 
 from mlagents.trainers.ppo.multi_gpu_policy import MultiGpuPPOPolicy

ml-agents/mlagents/trainers/tests/test_ppo.py

@@ -2,7 +2,8 @@
 import pytest
 
 import numpy as np
-import tensorflow as tf
+from mlagents.trainers import tf
+
 import yaml
 
 from mlagents.trainers.ppo.models import PPOModel

ml-agents/mlagents/trainers/tests/test_sac.py

@@ -3,7 +3,8 @@
 import yaml
 
 import numpy as np
-import tensorflow as tf
+from mlagents.trainers import tf
+
 
 from mlagents.trainers.sac.models import SACModel
 from mlagents.trainers.sac.policy import SACPolicy

ml-agents/mlagents/trainers/tests/test_trainer_controller.py

@@ -1,5 +1,7 @@
 from unittest.mock import MagicMock, Mock, patch
 
+from mlagents.trainers 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

ml-agents/mlagents/trainers/tf.py

@@ -0,0 +1,26 @@
+# 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")
+
+# A few things that we use live in different places between tensorflow 1.x and 2.x
+# If anything new is added, please add it here
+
+if _is_tensorflow2:
+    import tensorflow.compat.v1 as tf
+
+    tf_variance_scaling = tf.initializers.variance_scaling
+    tf_flatten = tf.layers.flatten
+    tf_rnn = tf.nn.rnn_cell
+
+    tf.disable_v2_behavior()
+else:
+    import tensorflow.contrib.layers as c_layers
+
+    tf_variance_scaling = c_layers.variance_scaling_initializer
+    tf_flatten = c_layers.flatten
+    tf_rnn = tf.contrib.rnn

ml-agents/mlagents/trainers/tf_policy.py

@@ -2,7 +2,7 @@
 from typing import Any, Dict
 
 import numpy as np
-import tensorflow as tf
+from mlagents.trainers import tf
 
 from mlagents.envs.exception import UnityException
 from mlagents.envs.policy import Policy

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.trainers import tf
+
 import numpy as np
 from collections import deque, defaultdict
 

ml-agents/mlagents/trainers/trainer_controller.py

@@ -8,7 +8,7 @@
 from typing import Dict, List, Optional
 
 import numpy as np
-import tensorflow as tf
+from mlagents.trainers import tf
 from time import time
 
 from mlagents.envs.env_manager import EnvironmentStep

ml-agents/setup.py

@@ -40,7 +40,7 @@
         "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",

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 instead (it handles tf2 compat).