hill-a · Miffyli · Oct 20, 2020 · Oct 20, 2020
diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst
@@ -16,6 +16,7 @@ Breaking Changes:
 New Features:
 ^^^^^^^^^^^^^
 - EvalCallback now works also for recurrent policies (@mily20001)
+- Add minimal support for TF2 using tensorflow.compat.v1 while keeping support for TF1
 
 Bug Fixes:
 ^^^^^^^^^^

diff --git a/stable_baselines/a2c/a2c.py b/stable_baselines/a2c/a2c.py
@@ -2,7 +2,7 @@
 
 import gym
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from stable_baselines import logger
 from stable_baselines.common import explained_variance, tf_util, ActorCriticRLModel, SetVerbosity, TensorboardWriter

diff --git a/stable_baselines/acer/acer_simple.py b/stable_baselines/acer/acer_simple.py
@@ -2,7 +2,7 @@
 import warnings
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from gym.spaces import Discrete, Box
 from collections import deque
 

diff --git a/stable_baselines/acktr/acktr.py b/stable_baselines/acktr/acktr.py
@@ -1,7 +1,7 @@
 import time
 import warnings
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from gym.spaces import Box, Discrete
 
 from stable_baselines import logger

diff --git a/stable_baselines/acktr/kfac.py b/stable_baselines/acktr/kfac.py
@@ -1,7 +1,7 @@
 import re
 from functools import reduce
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import numpy as np
 
 from stable_baselines.acktr.kfac_utils import detect_min_val, factor_reshape, gmatmul

diff --git a/stable_baselines/acktr/kfac_utils.py b/stable_baselines/acktr/kfac_utils.py
@@ -1,4 +1,4 @@
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def gmatmul(tensor_a, tensor_b, transpose_a=False, transpose_b=False, reduce_dim=None):

diff --git a/stable_baselines/common/base_class.py b/stable_baselines/common/base_class.py
@@ -10,7 +10,7 @@
 import gym
 import cloudpickle
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from stable_baselines.common.misc_util import set_global_seeds
 from stable_baselines.common.save_util import data_to_json, json_to_data, params_to_bytes, bytes_to_params

diff --git a/stable_baselines/common/distributions.py b/stable_baselines/common/distributions.py
@@ -1,5 +1,5 @@
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from tensorflow.python.ops import math_ops
 from gym import spaces
 

diff --git a/stable_baselines/common/input.py b/stable_baselines/common/input.py
@@ -1,5 +1,5 @@
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from gym.spaces import Discrete, Box, MultiBinary, MultiDiscrete
 
 

diff --git a/stable_baselines/common/misc_util.py b/stable_baselines/common/misc_util.py
@@ -2,7 +2,7 @@
 
 import gym
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def zipsame(*seqs):

diff --git a/stable_baselines/common/mpi_adam.py b/stable_baselines/common/mpi_adam.py
@@ -1,4 +1,4 @@
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import numpy as np
 import mpi4py
 

diff --git a/stable_baselines/common/mpi_running_mean_std.py b/stable_baselines/common/mpi_running_mean_std.py
@@ -1,5 +1,5 @@
 import mpi4py
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import numpy as np
 
 import stable_baselines.common.tf_util as tf_util

diff --git a/stable_baselines/common/policies.py b/stable_baselines/common/policies.py
@@ -3,7 +3,7 @@
 from abc import ABC, abstractmethod
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from gym.spaces import Discrete
 
 from stable_baselines.common.tf_util import batch_to_seq, seq_to_batch

diff --git a/stable_baselines/common/tf_layers.py b/stable_baselines/common/tf_layers.py
@@ -1,5 +1,5 @@
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def ortho_init(scale=1.0):

diff --git a/stable_baselines/common/tf_util.py b/stable_baselines/common/tf_util.py
@@ -5,7 +5,7 @@
 from typing import Set
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def is_image(tensor):

diff --git a/stable_baselines/ddpg/ddpg.py b/stable_baselines/ddpg/ddpg.py
@@ -7,7 +7,7 @@
 
 import gym
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import tensorflow.contrib as tc
 from mpi4py import MPI
 

diff --git a/stable_baselines/ddpg/main.py b/stable_baselines/ddpg/main.py
@@ -3,7 +3,7 @@
 import os
 
 import gym
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import numpy as np
 from mpi4py import MPI
 

diff --git a/stable_baselines/ddpg/policies.py b/stable_baselines/ddpg/policies.py
@@ -1,4 +1,4 @@
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from gym.spaces import Box
 
 from stable_baselines.common.policies import BasePolicy, nature_cnn, register_policy

diff --git a/stable_baselines/deepq/build_graph.py b/stable_baselines/deepq/build_graph.py
@@ -62,7 +62,7 @@
     Q' is set to Q once every 10000 updates training steps.
 
 """
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from gym.spaces import MultiDiscrete
 
 from stable_baselines.common import tf_util

diff --git a/stable_baselines/deepq/dqn.py b/stable_baselines/deepq/dqn.py
@@ -1,6 +1,6 @@
 from functools import partial
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import numpy as np
 import gym
 

diff --git a/stable_baselines/deepq/policies.py b/stable_baselines/deepq/policies.py
@@ -1,4 +1,4 @@
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import tensorflow.contrib.layers as tf_layers
 import numpy as np
 from gym.spaces import Discrete

diff --git a/stable_baselines/gail/adversary.py b/stable_baselines/gail/adversary.py
@@ -3,7 +3,7 @@
 I follow the architecture from the official repository
 """
 import gym
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import numpy as np
 
 from stable_baselines.common.mpi_running_mean_std import RunningMeanStd

diff --git a/stable_baselines/logger.py b/stable_baselines/logger.py
@@ -9,7 +9,7 @@
 from collections import defaultdict
 from typing import Optional
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from tensorflow.python import pywrap_tensorflow
 from tensorflow.core.util import event_pb2
 from tensorflow.python.util import compat
@@ -715,7 +715,7 @@ def read_tb(path):
     import numpy as np
     from glob import glob
     # from collections import defaultdict
-    import tensorflow as tf
+    import tensorflow.compat.v1 as tf
     if os.path.isdir(path):
         fnames = glob(os.path.join(path, "events.*"))
     elif os.path.basename(path).startswith("events."):

diff --git a/stable_baselines/ppo1/pposgd_simple.py b/stable_baselines/ppo1/pposgd_simple.py
@@ -3,7 +3,7 @@
 
 import gym
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from mpi4py import MPI
 
 from stable_baselines.common import Dataset, explained_variance, fmt_row, zipsame, ActorCriticRLModel, SetVerbosity, \

diff --git a/stable_baselines/ppo2/ppo2.py b/stable_baselines/ppo2/ppo2.py
@@ -2,7 +2,7 @@
 
 import gym
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from stable_baselines import logger
 from stable_baselines.common import explained_variance, ActorCriticRLModel, tf_util, SetVerbosity, TensorboardWriter

diff --git a/stable_baselines/sac/policies.py b/stable_baselines/sac/policies.py
@@ -1,4 +1,4 @@
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import numpy as np
 from gym.spaces import Box
 

diff --git a/stable_baselines/sac/sac.py b/stable_baselines/sac/sac.py
@@ -2,7 +2,7 @@
 import warnings
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from stable_baselines.common import tf_util, OffPolicyRLModel, SetVerbosity, TensorboardWriter
 from stable_baselines.common.vec_env import VecEnv

diff --git a/stable_baselines/td3/policies.py b/stable_baselines/td3/policies.py
@@ -1,4 +1,4 @@
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from gym.spaces import Box
 
 from stable_baselines.common.policies import BasePolicy, nature_cnn, register_policy

diff --git a/stable_baselines/td3/td3.py b/stable_baselines/td3/td3.py
@@ -2,7 +2,7 @@
 import warnings
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from stable_baselines import logger
 from stable_baselines.common import tf_util, OffPolicyRLModel, SetVerbosity, TensorboardWriter

diff --git a/stable_baselines/trpo_mpi/trpo_mpi.py b/stable_baselines/trpo_mpi/trpo_mpi.py
@@ -4,7 +4,7 @@
 
 import gym
 from mpi4py import MPI
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import numpy as np
 
 import stable_baselines.common.tf_util as tf_util

diff --git a/tests/test_a2c_conv.py b/tests/test_a2c_conv.py
@@ -1,6 +1,6 @@
 import gym
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from stable_baselines.common.tf_layers import conv
 from stable_baselines.common.input import observation_input

diff --git a/tests/test_custom_policy.py b/tests/test_custom_policy.py
@@ -2,7 +2,7 @@
 
 import gym
 import pytest
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from stable_baselines import A2C, ACER, ACKTR, DQN, PPO1, PPO2, TRPO, SAC, DDPG
 from stable_baselines.common.policies import FeedForwardPolicy

diff --git a/tests/test_distri.py b/tests/test_distri.py
@@ -1,5 +1,5 @@
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 import stable_baselines.common.tf_util as tf_util
 from stable_baselines.common.distributions import DiagGaussianProbabilityDistributionType,\

diff --git a/tests/test_math_util.py b/tests/test_math_util.py
@@ -1,4 +1,4 @@
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import numpy as np
 from gym.spaces.box import Box
 

diff --git a/tests/test_tf_util.py b/tests/test_tf_util.py
@@ -1,6 +1,6 @@
 # tests for tf_util
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from stable_baselines.common.tf_util import function, initialize, single_threaded_session, is_image