[refactor] Move TF-specific files to tf/ folder

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

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
 
 

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
 

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.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.distributions import (
+from mlagents.trainers.tf.distributions import (
     GaussianDistribution,
     MultiCategoricalDistribution,
 )

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

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.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

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

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.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

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

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,
 )

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
 

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

@@ -9,10 +9,10 @@
 
 
 from mlagents.trainers.policy.tf_policy import TFPolicy
-from mlagents.trainers.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__
 

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.models import EncoderType, ScheduleType
 from mlagents_envs.side_channel.environment_parameters_channel import (
     EnvironmentParametersChannel,
 )

ml-agents/mlagents/trainers/distributions.py → 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
 

ml-agents/mlagents/trainers/models.py → 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

ml-agents/mlagents/trainers/tensorflow_to_barracuda.py → 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)