From 577c4d71c5795f3153aaf23064cb0b62bf95de2b Mon Sep 17 00:00:00 2001
From: Ervin Teng <ervin@unity3d.com>
Date: Thu, 23 Jul 2020 15:01:01 -0700
Subject: [PATCH 1/2] Move TF-specific files to tf/ folder

---
 ml-agents/mlagents/model_serialization.py                     | 2 +-
 .../trainers/components/reward_signals/curiosity/model.py     | 2 +-
 .../mlagents/trainers/components/reward_signals/gail/model.py | 2 +-
 ml-agents/mlagents/trainers/policy/tf_policy.py               | 4 ++--
 ml-agents/mlagents/trainers/ppo/optimizer.py                  | 2 +-
 ml-agents/mlagents/trainers/sac/network.py                    | 2 +-
 ml-agents/mlagents/trainers/sac/optimizer.py                  | 2 +-
 ml-agents/mlagents/trainers/settings.py                       | 2 +-
 ml-agents/mlagents/trainers/tests/test_barracuda_converter.py | 2 +-
 ml-agents/mlagents/trainers/tests/test_distributions.py       | 2 +-
 ml-agents/mlagents/trainers/tests/test_models.py              | 2 +-
 ml-agents/mlagents/trainers/tests/test_nn_policy.py           | 2 +-
 ml-agents/mlagents/trainers/tests/test_simple_rl.py           | 2 +-
 ml-agents/mlagents/trainers/tf/__init__.py                    | 0
 ml-agents/mlagents/trainers/{ => tf}/distributions.py         | 2 +-
 ml-agents/mlagents/trainers/{ => tf}/models.py                | 0
 .../mlagents/trainers/{ => tf}/tensorflow_to_barracuda.py     | 2 +-
 17 files changed, 16 insertions(+), 16 deletions(-)
 create mode 100644 ml-agents/mlagents/trainers/tf/__init__.py
 rename ml-agents/mlagents/trainers/{ => tf}/distributions.py (99%)
 rename ml-agents/mlagents/trainers/{ => tf}/models.py (100%)
 rename ml-agents/mlagents/trainers/{ => tf}/tensorflow_to_barracuda.py (99%)

diff --git a/ml-agents/mlagents/model_serialization.py b/ml-agents/mlagents/model_serialization.py
index fedae58080..edc7a5f6ee 100644
--- a/ml-agents/mlagents/model_serialization.py
+++ b/ml-agents/mlagents/model_serialization.py
@@ -19,7 +19,7 @@
 from tensorflow.python.framework import graph_util
 
 from mlagents_envs.logging_util import get_logger
-from mlagents.trainers import tensorflow_to_barracuda as tf2bc
+from mlagents.trainers.tf import tensorflow_to_barracuda as tf2bc
 
 if LooseVersion(tf.__version__) < LooseVersion("1.12.0"):
     # ONNX is only tested on 1.12.0 and later
diff --git a/ml-agents/mlagents/trainers/components/reward_signals/curiosity/model.py b/ml-agents/mlagents/trainers/components/reward_signals/curiosity/model.py
index 684d111094..906ac549a4 100644
--- a/ml-agents/mlagents/trainers/components/reward_signals/curiosity/model.py
+++ b/ml-agents/mlagents/trainers/components/reward_signals/curiosity/model.py
@@ -1,7 +1,7 @@
 from typing import List, Tuple
 from mlagents.tf_utils import tf
 
-from mlagents.trainers.models import ModelUtils
+from mlagents.trainers.tf.models import ModelUtils
 from mlagents.trainers.policy.tf_policy import TFPolicy
 
 
diff --git a/ml-agents/mlagents/trainers/components/reward_signals/gail/model.py b/ml-agents/mlagents/trainers/components/reward_signals/gail/model.py
index fcb0a11d6b..78bd2dd7b5 100644
--- a/ml-agents/mlagents/trainers/components/reward_signals/gail/model.py
+++ b/ml-agents/mlagents/trainers/components/reward_signals/gail/model.py
@@ -3,7 +3,7 @@
 from mlagents.tf_utils import tf
 
 from mlagents.trainers.policy.tf_policy import TFPolicy
-from mlagents.trainers.models import ModelUtils
+from mlagents.trainers.tf.models import ModelUtils
 
 EPSILON = 1e-7
 
diff --git a/ml-agents/mlagents/trainers/policy/tf_policy.py b/ml-agents/mlagents/trainers/policy/tf_policy.py
index 261d7bf7c9..ec5fc7ce32 100644
--- a/ml-agents/mlagents/trainers/policy/tf_policy.py
+++ b/ml-agents/mlagents/trainers/policy/tf_policy.py
@@ -15,10 +15,10 @@
 from mlagents.trainers.trajectory import SplitObservations
 from mlagents.trainers.behavior_id_utils import get_global_agent_id
 from mlagents_envs.base_env import DecisionSteps
-from mlagents.trainers.models import ModelUtils, EncoderType
+from mlagents.trainers.tf.models import ModelUtils, EncoderType
 from mlagents.trainers.settings import TrainerSettings
 from mlagents.trainers import __version__
-from mlagents.trainers.distributions import (
+from mlagents.trainers.tf.distributions import (
     GaussianDistribution,
     MultiCategoricalDistribution,
 )
diff --git a/ml-agents/mlagents/trainers/ppo/optimizer.py b/ml-agents/mlagents/trainers/ppo/optimizer.py
index 0df43eeb1f..7cbeb20c0a 100644
--- a/ml-agents/mlagents/trainers/ppo/optimizer.py
+++ b/ml-agents/mlagents/trainers/ppo/optimizer.py
@@ -2,7 +2,7 @@
 import numpy as np
 from mlagents.tf_utils import tf
 from mlagents_envs.timers import timed
-from mlagents.trainers.models import ModelUtils, EncoderType
+from mlagents.trainers.tf.models import ModelUtils, EncoderType
 from mlagents.trainers.policy.tf_policy import TFPolicy
 from mlagents.trainers.optimizer.tf_optimizer import TFOptimizer
 from mlagents.trainers.buffer import AgentBuffer
diff --git a/ml-agents/mlagents/trainers/sac/network.py b/ml-agents/mlagents/trainers/sac/network.py
index 5c0e8fa0e8..99e397f9cf 100644
--- a/ml-agents/mlagents/trainers/sac/network.py
+++ b/ml-agents/mlagents/trainers/sac/network.py
@@ -1,6 +1,6 @@
 from typing import Dict, Optional
 from mlagents.tf_utils import tf
-from mlagents.trainers.models import ModelUtils, EncoderType
+from mlagents.trainers.tf.models import ModelUtils, EncoderType
 
 LOG_STD_MAX = 2
 LOG_STD_MIN = -20
diff --git a/ml-agents/mlagents/trainers/sac/optimizer.py b/ml-agents/mlagents/trainers/sac/optimizer.py
index b629dc153f..9c5d0e6b83 100644
--- a/ml-agents/mlagents/trainers/sac/optimizer.py
+++ b/ml-agents/mlagents/trainers/sac/optimizer.py
@@ -5,7 +5,7 @@
 
 from mlagents_envs.logging_util import get_logger
 from mlagents.trainers.sac.network import SACPolicyNetwork, SACTargetNetwork
-from mlagents.trainers.models import ModelUtils
+from mlagents.trainers.tf.models import ModelUtils
 from mlagents.trainers.optimizer.tf_optimizer import TFOptimizer
 from mlagents.trainers.policy.tf_policy import TFPolicy
 from mlagents.trainers.buffer import AgentBuffer
diff --git a/ml-agents/mlagents/trainers/settings.py b/ml-agents/mlagents/trainers/settings.py
index 9b75496ffc..3a20f9c3d1 100644
--- a/ml-agents/mlagents/trainers/settings.py
+++ b/ml-agents/mlagents/trainers/settings.py
@@ -11,7 +11,7 @@
 from mlagents.trainers.cli_utils import StoreConfigFile, DetectDefault, parser
 from mlagents.trainers.cli_utils import load_config
 from mlagents.trainers.exception import TrainerConfigError
-from mlagents.trainers.models import ScheduleType, EncoderType
+from mlagents.trainers.tf.models import ScheduleType, EncoderType
 
 from mlagents_envs import logging_util
 from mlagents_envs.side_channel.environment_parameters_channel import (
diff --git a/ml-agents/mlagents/trainers/tests/test_barracuda_converter.py b/ml-agents/mlagents/trainers/tests/test_barracuda_converter.py
index b1d504e56a..a01cd36e35 100644
--- a/ml-agents/mlagents/trainers/tests/test_barracuda_converter.py
+++ b/ml-agents/mlagents/trainers/tests/test_barracuda_converter.py
@@ -2,7 +2,7 @@
 import tempfile
 import pytest
 
-import mlagents.trainers.tensorflow_to_barracuda as tf2bc
+import mlagents.trainers.tf.tensorflow_to_barracuda as tf2bc
 from mlagents.trainers.tests.test_nn_policy import create_policy_mock
 from mlagents.trainers.settings import TrainerSettings
 from mlagents.tf_utils import tf
diff --git a/ml-agents/mlagents/trainers/tests/test_distributions.py b/ml-agents/mlagents/trainers/tests/test_distributions.py
index 30756f12a2..c098d067f0 100644
--- a/ml-agents/mlagents/trainers/tests/test_distributions.py
+++ b/ml-agents/mlagents/trainers/tests/test_distributions.py
@@ -2,7 +2,7 @@
 
 from mlagents.tf_utils import tf
 
-from mlagents.trainers.distributions import (
+from mlagents.trainers.tf.distributions import (
     GaussianDistribution,
     MultiCategoricalDistribution,
 )
diff --git a/ml-agents/mlagents/trainers/tests/test_models.py b/ml-agents/mlagents/trainers/tests/test_models.py
index 54ce8f3f55..1736d89315 100644
--- a/ml-agents/mlagents/trainers/tests/test_models.py
+++ b/ml-agents/mlagents/trainers/tests/test_models.py
@@ -1,6 +1,6 @@
 import pytest
 
-from mlagents.trainers.models import ModelUtils
+from mlagents.trainers.tf.models import ModelUtils
 from mlagents.tf_utils import tf
 from mlagents_envs.base_env import BehaviorSpec, ActionType
 
diff --git a/ml-agents/mlagents/trainers/tests/test_nn_policy.py b/ml-agents/mlagents/trainers/tests/test_nn_policy.py
index a09d1a3b23..93a373e2c2 100644
--- a/ml-agents/mlagents/trainers/tests/test_nn_policy.py
+++ b/ml-agents/mlagents/trainers/tests/test_nn_policy.py
@@ -9,7 +9,7 @@
 
 
 from mlagents.trainers.policy.tf_policy import TFPolicy
-from mlagents.trainers.models import EncoderType, ModelUtils, Tensor3DShape
+from mlagents.trainers.tf.models import EncoderType, ModelUtils, Tensor3DShape
 from mlagents.trainers.exception import UnityTrainerException
 from mlagents.trainers.tests import mock_brain as mb
 from mlagents.trainers.settings import TrainerSettings, NetworkSettings
diff --git a/ml-agents/mlagents/trainers/tests/test_simple_rl.py b/ml-agents/mlagents/trainers/tests/test_simple_rl.py
index bf75a3b9b4..c348ae2921 100644
--- a/ml-agents/mlagents/trainers/tests/test_simple_rl.py
+++ b/ml-agents/mlagents/trainers/tests/test_simple_rl.py
@@ -27,7 +27,7 @@
     RewardSignalType,
 )
 from mlagents.trainers.environment_parameter_manager import EnvironmentParameterManager
-from mlagents.trainers.models import EncoderType, ScheduleType
+from mlagents.trainers.tf.models import EncoderType, ScheduleType
 from mlagents_envs.side_channel.environment_parameters_channel import (
     EnvironmentParametersChannel,
 )
diff --git a/ml-agents/mlagents/trainers/tf/__init__.py b/ml-agents/mlagents/trainers/tf/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/ml-agents/mlagents/trainers/distributions.py b/ml-agents/mlagents/trainers/tf/distributions.py
similarity index 99%
rename from ml-agents/mlagents/trainers/distributions.py
rename to ml-agents/mlagents/trainers/tf/distributions.py
index 294bad11cb..1d641d86a7 100644
--- a/ml-agents/mlagents/trainers/distributions.py
+++ b/ml-agents/mlagents/trainers/tf/distributions.py
@@ -3,7 +3,7 @@
 import numpy as np
 
 from mlagents.tf_utils import tf
-from mlagents.trainers.models import ModelUtils
+from mlagents.trainers.tf.models import ModelUtils
 
 EPSILON = 1e-6  # Small value to avoid divide by zero
 
diff --git a/ml-agents/mlagents/trainers/models.py b/ml-agents/mlagents/trainers/tf/models.py
similarity index 100%
rename from ml-agents/mlagents/trainers/models.py
rename to ml-agents/mlagents/trainers/tf/models.py
diff --git a/ml-agents/mlagents/trainers/tensorflow_to_barracuda.py b/ml-agents/mlagents/trainers/tf/tensorflow_to_barracuda.py
similarity index 99%
rename from ml-agents/mlagents/trainers/tensorflow_to_barracuda.py
rename to ml-agents/mlagents/trainers/tf/tensorflow_to_barracuda.py
index 0958b92a3a..80d5371812 100644
--- a/ml-agents/mlagents/trainers/tensorflow_to_barracuda.py
+++ b/ml-agents/mlagents/trainers/tf/tensorflow_to_barracuda.py
@@ -23,7 +23,7 @@
     )
     # Te following code can be used as an example of API used from another module
     # convert() is the main entry point for converter
-    import tensorflow_to_barracuda as tf2bc
+    import tf.tensorflow_to_barracuda as tf2bc
 
     tf2bc.convert(args.source_file, args.target_file, args.trim_unused_by_output, args)
 

From f6d5ae5d77b194f6d133cda7b3b461ed27be8472 Mon Sep 17 00:00:00 2001
From: Ervin Teng <ervin@unity3d.com>
Date: Thu, 23 Jul 2020 15:14:21 -0700
Subject: [PATCH 2/2] Move EncoderType and ScheduleType to settings.py

---
 ml-agents/mlagents/trainers/policy/tf_policy.py     |  4 ++--
 ml-agents/mlagents/trainers/sac/network.py          |  3 ++-
 ml-agents/mlagents/trainers/settings.py             | 12 +++++++++++-
 ml-agents/mlagents/trainers/tests/test_nn_policy.py |  4 ++--
 ml-agents/mlagents/trainers/tests/test_simple_rl.py |  3 ++-
 ml-agents/mlagents/trainers/tf/models.py            | 13 +------------
 6 files changed, 20 insertions(+), 19 deletions(-)

diff --git a/ml-agents/mlagents/trainers/policy/tf_policy.py b/ml-agents/mlagents/trainers/policy/tf_policy.py
index ec5fc7ce32..da6580c0a8 100644
--- a/ml-agents/mlagents/trainers/policy/tf_policy.py
+++ b/ml-agents/mlagents/trainers/policy/tf_policy.py
@@ -15,8 +15,8 @@
 from mlagents.trainers.trajectory import SplitObservations
 from mlagents.trainers.behavior_id_utils import get_global_agent_id
 from mlagents_envs.base_env import DecisionSteps
-from mlagents.trainers.tf.models import ModelUtils, EncoderType
-from mlagents.trainers.settings import TrainerSettings
+from mlagents.trainers.tf.models import ModelUtils
+from mlagents.trainers.settings import TrainerSettings, EncoderType
 from mlagents.trainers import __version__
 from mlagents.trainers.tf.distributions import (
     GaussianDistribution,
diff --git a/ml-agents/mlagents/trainers/sac/network.py b/ml-agents/mlagents/trainers/sac/network.py
index 99e397f9cf..568de44753 100644
--- a/ml-agents/mlagents/trainers/sac/network.py
+++ b/ml-agents/mlagents/trainers/sac/network.py
@@ -1,6 +1,7 @@
 from typing import Dict, Optional
 from mlagents.tf_utils import tf
-from mlagents.trainers.tf.models import ModelUtils, EncoderType
+from mlagents.trainers.tf.models import ModelUtils
+from mlagents.trainers.settings import EncoderType
 
 LOG_STD_MAX = 2
 LOG_STD_MIN = -20
diff --git a/ml-agents/mlagents/trainers/settings.py b/ml-agents/mlagents/trainers/settings.py
index 3a20f9c3d1..89ee40bc69 100644
--- a/ml-agents/mlagents/trainers/settings.py
+++ b/ml-agents/mlagents/trainers/settings.py
@@ -11,7 +11,6 @@
 from mlagents.trainers.cli_utils import StoreConfigFile, DetectDefault, parser
 from mlagents.trainers.cli_utils import load_config
 from mlagents.trainers.exception import TrainerConfigError
-from mlagents.trainers.tf.models import ScheduleType, EncoderType
 
 from mlagents_envs import logging_util
 from mlagents_envs.side_channel.environment_parameters_channel import (
@@ -51,6 +50,17 @@ def as_dict(self):
         return cattr.unstructure(self)
 
 
+class EncoderType(Enum):
+    SIMPLE = "simple"
+    NATURE_CNN = "nature_cnn"
+    RESNET = "resnet"
+
+
+class ScheduleType(Enum):
+    CONSTANT = "constant"
+    LINEAR = "linear"
+
+
 @attr.s(auto_attribs=True)
 class NetworkSettings:
     @attr.s
diff --git a/ml-agents/mlagents/trainers/tests/test_nn_policy.py b/ml-agents/mlagents/trainers/tests/test_nn_policy.py
index 93a373e2c2..3492de2d45 100644
--- a/ml-agents/mlagents/trainers/tests/test_nn_policy.py
+++ b/ml-agents/mlagents/trainers/tests/test_nn_policy.py
@@ -9,10 +9,10 @@
 
 
 from mlagents.trainers.policy.tf_policy import TFPolicy
-from mlagents.trainers.tf.models import EncoderType, ModelUtils, Tensor3DShape
+from mlagents.trainers.tf.models import ModelUtils, Tensor3DShape
 from mlagents.trainers.exception import UnityTrainerException
 from mlagents.trainers.tests import mock_brain as mb
-from mlagents.trainers.settings import TrainerSettings, NetworkSettings
+from mlagents.trainers.settings import TrainerSettings, NetworkSettings, EncoderType
 from mlagents.trainers.tests.test_trajectory import make_fake_trajectory
 from mlagents.trainers import __version__
 
diff --git a/ml-agents/mlagents/trainers/tests/test_simple_rl.py b/ml-agents/mlagents/trainers/tests/test_simple_rl.py
index c348ae2921..c23f0ec254 100644
--- a/ml-agents/mlagents/trainers/tests/test_simple_rl.py
+++ b/ml-agents/mlagents/trainers/tests/test_simple_rl.py
@@ -25,9 +25,10 @@
     GAILSettings,
     TrainerType,
     RewardSignalType,
+    EncoderType,
+    ScheduleType,
 )
 from mlagents.trainers.environment_parameter_manager import EnvironmentParameterManager
-from mlagents.trainers.tf.models import EncoderType, ScheduleType
 from mlagents_envs.side_channel.environment_parameters_channel import (
     EnvironmentParametersChannel,
 )
diff --git a/ml-agents/mlagents/trainers/tf/models.py b/ml-agents/mlagents/trainers/tf/models.py
index 7fae0d9a6b..18ad6b0b30 100644
--- a/ml-agents/mlagents/trainers/tf/models.py
+++ b/ml-agents/mlagents/trainers/tf/models.py
@@ -1,8 +1,8 @@
-from enum import Enum
 from typing import Callable, Dict, List, Tuple, NamedTuple
 
 import numpy as np
 from mlagents.tf_utils import tf
+from mlagents.trainers.settings import EncoderType, ScheduleType
 
 from mlagents.trainers.exception import UnityTrainerException
 
@@ -20,17 +20,6 @@ class Tensor3DShape(NamedTuple):
     num_channels: int
 
 
-class EncoderType(Enum):
-    SIMPLE = "simple"
-    NATURE_CNN = "nature_cnn"
-    RESNET = "resnet"
-
-
-class ScheduleType(Enum):
-    CONSTANT = "constant"
-    LINEAR = "linear"
-
-
 class NormalizerTensors(NamedTuple):
     update_op: tf.Operation
     steps: tf.Tensor