centralize tensorflow imports

ml-agents/mlagents/trainers/__init__.py

-import logging
-
-def warnings_as_errors(log_record):
-    # Raise deprecated warnings as exceptions.
-    if log_record.levelno == logging.WARNING and "deprecated" in log_record.msg:
-        merged = log_record.getMessage()
-        raise RuntimeError(merged)
-    return True
-
-# TODO only enable this with a environment variable
-if False:
-    logging.getLogger('tensorflow').addFilter(warnings_as_errors)
-
-
-# TODO better place to put this? move everything to tf.py?
-try:
-    import tensorflow.compat.v1 as tf
-except ImportError:
-    import tensorflow as tf
-tf.disable_v2_behavior()
+from mlagents.trainers.tf import tf as tf

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

-try:
-    import tensorflow.compat.v1 as tf
-except ImportError:
-    import tensorflow as tf
+from mlagents.trainers import tf
-if True: # TODO TF2
+if True:  # TODO TF2
+
     tf_variance_scaling = c_layers.variance_scaling_initializer
     tf_flatten = c_layers.flatten
 
                         size,
                         activation=None,
                         use_bias=False,
-                        kernel_initializer=tf_variance_scaling(
-                            0.01
-                        ),
+                        kernel_initializer=tf_variance_scaling(0.01),
                     )
                 )
             self.action_probs = tf.concat(

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

-try:
-    import tensorflow.compat.v1 as tf
-except ImportError:
-    import tensorflow as tf
+from mlagents.trainers import tf
 
 from mlagents.trainers.models import LearningModel
 

ml-agents/mlagents/trainers/components/reward_signals/__init__.py

 import numpy as np
 import abc
 
-try:
-    import tensorflow.compat.v1 as tf
-except ImportError:
-    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

 from typing import List, Tuple
-try:
-    import tensorflow.compat.v1 as tf
-except ImportError:
-    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

 from typing import Any, Dict, List
 import numpy as np
-try:
-    import tensorflow.compat.v1 as tf
-except ImportError:
-    import tensorflow as tf
+from mlagents.trainers import tf
 
 from mlagents.envs.brain import BrainInfo
 

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

 from typing import Tuple, List
 
-try:
-    import tensorflow.compat.v1 as tf
-except ImportError:
-    import tensorflow as tf
+from mlagents.trainers import tf
 
 from mlagents.trainers.models import LearningModel
 

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

 from typing import Any, Dict, List
 import logging
 import numpy as np
-try:
-    import tensorflow.compat.v1 as tf
-except ImportError:
-    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

 from typing import Callable, List
 
 import numpy as np
-try:
-    import tensorflow.compat.v1 as tf
-except ImportError:
-    import tensorflow as tf
+from mlagents.trainers import tf
-if True: # TODO TF2
+if True:  # TODO TF2
+
     tf_variance_scaling = c_layers.variance_scaling_initializer
     tf_flatten = c_layers.flatten
     tf_rnn = tf.contrib.rnn

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

 import logging
 import numpy as np
 
-try:
-    import tensorflow.compat.v1 as tf
-except ImportError:
-    import tensorflow as tf
+from mlagents.trainers import tf
 
 from mlagents.trainers.models import LearningModel, EncoderType, LearningRateSchedule
 

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

 import logging
 
-try:
-    import tensorflow.compat.v1 as tf
-except ImportError:
-    import tensorflow as tf
+from mlagents.trainers import tf
 
 from tensorflow.python.client import device_lib
 from mlagents.envs.timers import timed

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

 import numpy as np
 from typing import Any, Dict, Optional
 
-try:
-    import tensorflow.compat.v1 as tf
-except ImportError:
-    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

 import logging
 import numpy as np
 
-try:
-    import tensorflow.compat.v1 as tf
-except ImportError:
-    import tensorflow as tf
+from mlagents.trainers import tf
-if True: # TODO TF2
+if True:  # TODO TF2
+
     tf_variance_scaling = c_layers.variance_scaling_initializer
     tf_flatten = c_layers.flatten
 
                         size,
                         activation=None,
                         use_bias=False,
-                        kernel_initializer=tf_variance_scaling(
-                            0.01
-                        ),
+                        kernel_initializer=tf_variance_scaling(0.01),
                     )
                 )
             all_logits = tf.concat(

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

 import logging
 from typing import Dict, Any, Optional
 import numpy as np
-try:
-    import tensorflow.compat.v1 as tf
-except ImportError:
-    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

 from __future__ import print_function
 import numpy as np
 import struct  # convert from Python values and C structs
-try:
-    import tensorflow.compat.v1 as tf
-except ImportError:
-    import tensorflow as tf
+from mlagents.trainers import tf
 import re
 
 # import barracuda

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

 import os
 
 import numpy as np
-try:
-    import tensorflow.compat.v1 as tf
-except ImportError:
-    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

 import unittest.mock as mock
 import pytest
 
-try:
-    import tensorflow.compat.v1 as tf
-except ImportError:
-    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

 import pytest
 
 import numpy as np
-try:
-    import tensorflow.compat.v1 as tf
-except ImportError:
-    import tensorflow as tf
+from mlagents.trainers import tf
 
 import yaml
 

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

 import yaml
 
 import numpy as np
-try:
-    import tensorflow.compat.v1 as tf
-except ImportError:
-    import tensorflow as tf
+from mlagents.trainers import tf
 
 
 from mlagents.trainers.sac.models import SACModel

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

 from unittest.mock import MagicMock, Mock, patch
 
-try:
-    import tensorflow.compat.v1 as tf
-except ImportError:
-    import tensorflow as tf
+from mlagents.trainers import tf
 
 import yaml
 import pytest

ml-agents/mlagents/trainers/tf.py

+import logging
+
+
+def warnings_as_errors(log_record):
+    # Raise deprecated warnings as exceptions.
+    if log_record.levelno == logging.WARNING and "deprecated" in log_record.msg:
+        merged = log_record.getMessage()
+        raise RuntimeError(merged)
+    return True
+
+
+# TODO only enable this with a environment variable
+if False:
+    logging.getLogger("tensorflow").addFilter(warnings_as_errors)
+
+try:
+    import tensorflow.compat.v1 as tf
+except ImportError:
+    import tensorflow as tf
+
+is_tf2 = True
+if is_tf2:
+    tf.disable_v2_behavior()

ml-agents/mlagents/trainers/tf_policy.py

 from typing import Any, Dict
 
 import numpy as np
-try:
-    import tensorflow.compat.v1 as tf
-except ImportError:
-    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

 from typing import Dict, List, Deque, Any
 import os
 
-try:
-    import tensorflow.compat.v1 as tf
-except ImportError:
-    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

 from typing import Dict, List, Optional
 
 import numpy as np
-try:
-    import tensorflow.compat.v1 as tf
-except ImportError:
-    import tensorflow as tf
+from mlagents.trainers import tf
 from time import time
 
 from mlagents.envs.env_manager import EnvironmentStep