[change] Organize trainer files a bit better (#3538)

ml-agents/mlagents/trainers/agent_processor.py

 
 from mlagents_envs.base_env import BatchedStepResult, StepResult
 from mlagents.trainers.trajectory import Trajectory, AgentExperience
-from mlagents.trainers.tf_policy import TFPolicy
+from mlagents.trainers.policy.tf_policy import TFPolicy
 from mlagents.trainers.policy import Policy
 from mlagents.trainers.action_info import ActionInfo, ActionInfoOutputs
 from mlagents.trainers.stats import StatsReporter

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

 from mlagents.tf_utils import tf
 
-from mlagents.trainers.tf_policy import TFPolicy
+from mlagents.trainers.policy.tf_policy import TFPolicy
 
 
 class BCModel(object):

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

 from typing import Dict, Any
 import numpy as np
 
-from mlagents.trainers.tf_policy import TFPolicy
+from mlagents.trainers.policy.tf_policy import TFPolicy
 from .model import BCModel
 from mlagents.trainers.demo_loader import demo_to_buffer
 from mlagents.trainers.exception import UnityTrainerException

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

 from mlagents.tf_utils import tf
 
 from mlagents.trainers.exception import UnityTrainerException
-from mlagents.trainers.tf_policy import TFPolicy
+from mlagents.trainers.policy.tf_policy import TFPolicy
 
 logger = logging.getLogger("mlagents.trainers")
 

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

 from mlagents.tf_utils import tf
 
 from mlagents.trainers.models import ModelUtils
-from mlagents.trainers.tf_policy import TFPolicy
+from mlagents.trainers.policy.tf_policy import TFPolicy
 
 
 class CuriosityModel(object):

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

 
 from mlagents.trainers.components.reward_signals import RewardSignal, RewardSignalResult
 from mlagents.trainers.components.reward_signals.curiosity.model import CuriosityModel
-from mlagents.trainers.tf_policy import TFPolicy
+from mlagents.trainers.policy.tf_policy import TFPolicy
 
 
 class CuriosityRewardSignal(RewardSignal):

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

 
 from mlagents.tf_utils import tf
 
-from mlagents.trainers.tf_policy import TFPolicy
+from mlagents.trainers.policy.tf_policy import TFPolicy
 from mlagents.trainers.models import ModelUtils
 
 EPSILON = 1e-7

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

 from mlagents.tf_utils import tf
 
 from mlagents.trainers.components.reward_signals import RewardSignal, RewardSignalResult
-from mlagents.trainers.tf_policy import TFPolicy
+from mlagents.trainers.policy.tf_policy import TFPolicy
 from .model import GAILModel
 from mlagents.trainers.demo_loader import demo_to_buffer
 

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

 from mlagents.trainers.components.reward_signals.curiosity.signal import (
     CuriosityRewardSignal,
 )
-from mlagents.trainers.tf_policy import TFPolicy
+from mlagents.trainers.policy.tf_policy import TFPolicy
 
 logger = logging.getLogger("mlagents.trainers")
 

ml-agents/mlagents/trainers/env_manager.py

 from typing import List, Dict, NamedTuple, Iterable
 from mlagents_envs.base_env import BatchedStepResult, AgentGroupSpec, AgentGroup
 from mlagents.trainers.brain import BrainParameters
-from mlagents.trainers.tf_policy import TFPolicy
+from mlagents.trainers.policy.tf_policy import TFPolicy
 from mlagents.trainers.agent_processor import AgentManager, AgentManagerQueue
 from mlagents.trainers.action_info import ActionInfo
 

ml-agents/mlagents/trainers/ghost/trainer.py

 
 from mlagents.trainers.brain import BrainParameters
 from mlagents.trainers.policy import Policy
-from mlagents.trainers.tf_policy import TFPolicy
+from mlagents.trainers.policy.tf_policy import TFPolicy
 
 from mlagents.trainers.trainer import Trainer
 from mlagents.trainers.trajectory import Trajectory

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

 from mlagents.tf_utils import tf
 from mlagents_envs.timers import timed
 from mlagents.trainers.models import ModelUtils, EncoderType, LearningRateSchedule
-from mlagents.trainers.tf_policy import TFPolicy
-from mlagents.trainers.common.tf_optimizer import TFOptimizer
+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/ppo/trainer.py

 
 import numpy as np
 
-from mlagents.trainers.common.nn_policy import NNPolicy
-from mlagents.trainers.rl_trainer import RLTrainer
+from mlagents.trainers.policy.nn_policy import NNPolicy
+from mlagents.trainers.trainer.rl_trainer import RLTrainer
-from mlagents.trainers.tf_policy import TFPolicy
+from mlagents.trainers.policy.tf_policy import TFPolicy
 from mlagents.trainers.ppo.optimizer import PPOOptimizer
 from mlagents.trainers.trajectory import Trajectory
 

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

 
 from mlagents.trainers.sac.network import SACPolicyNetwork, SACTargetNetwork
 from mlagents.trainers.models import LearningRateSchedule, EncoderType, ModelUtils
-from mlagents.trainers.common.tf_optimizer import TFOptimizer
-from mlagents.trainers.tf_policy import TFPolicy
+from mlagents.trainers.optimizer.tf_optimizer import TFOptimizer
+from mlagents.trainers.policy.tf_policy import TFPolicy
 from mlagents.trainers.buffer import AgentBuffer
 from mlagents_envs.timers import timed
 

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

 
 
 from mlagents_envs.timers import timed
-from mlagents.trainers.tf_policy import TFPolicy
-from mlagents.trainers.common.nn_policy import NNPolicy
+from mlagents.trainers.policy.tf_policy import TFPolicy
+from mlagents.trainers.policy.nn_policy import NNPolicy
-from mlagents.trainers.rl_trainer import RLTrainer
+from mlagents.trainers.trainer.rl_trainer import RLTrainer
 from mlagents.trainers.trajectory import Trajectory, SplitObservations
 from mlagents.trainers.brain import BrainParameters
 

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

 import yaml
 import os
 
-from mlagents.trainers.common.nn_policy import NNPolicy
+from mlagents.trainers.policy.nn_policy import NNPolicy
 from mlagents.trainers.components.bc.module import BCModule
 
 

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

 
 import yaml
 
-from mlagents.trainers.common.distributions import (
+from mlagents.trainers.distributions import (
     GaussianDistribution,
     MultiCategoricalDistribution,
 )

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

 
 import yaml
 
-from mlagents.trainers.common.nn_policy import NNPolicy
+from mlagents.trainers.policy.nn_policy import NNPolicy
 from mlagents.trainers.models import EncoderType, ModelUtils
 from mlagents.trainers.exception import UnityTrainerException
 from mlagents.trainers.brain import BrainParameters, CameraResolution

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

-from mlagents.trainers.tf_policy import TFPolicy
+from mlagents.trainers.policy.tf_policy import TFPolicy
 from mlagents_envs.base_env import BatchedStepResult, AgentGroupSpec
 from mlagents.trainers.action_info import ActionInfo
 from unittest.mock import MagicMock

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

 
 from mlagents.trainers.ppo.trainer import PPOTrainer, discount_rewards
 from mlagents.trainers.ppo.optimizer import PPOOptimizer
-from mlagents.trainers.common.nn_policy import NNPolicy
+from mlagents.trainers.policy.nn_policy import NNPolicy
 from mlagents.trainers.brain import BrainParameters
 from mlagents.trainers.agent_processor import AgentManagerQueue
 from mlagents.trainers.tests import mock_brain as mb

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

 import yaml
 import os
 import mlagents.trainers.tests.mock_brain as mb
-from mlagents.trainers.common.nn_policy import NNPolicy
+from mlagents.trainers.policy.nn_policy import NNPolicy
 from mlagents.trainers.sac.optimizer import SACOptimizer
 from mlagents.trainers.ppo.optimizer import PPOOptimizer
 

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

 import yaml
 from unittest import mock
 import mlagents.trainers.tests.mock_brain as mb
-from mlagents.trainers.rl_trainer import RLTrainer
+from mlagents.trainers.trainer.rl_trainer import RLTrainer
 from mlagents.trainers.tests.test_buffer import construct_fake_buffer
 from mlagents.trainers.agent_processor import AgentManagerQueue
 
         assert len(arr) == 0
 
 
-@mock.patch("mlagents.trainers.rl_trainer.RLTrainer.clear_update_buffer")
+@mock.patch("mlagents.trainers.trainer.rl_trainer.RLTrainer.clear_update_buffer")
 def test_advance(mocked_clear_update_buffer):
     trainer = create_rl_trainer()
     trajectory_queue = AgentManagerQueue("testbrain")

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

 
 from mlagents.trainers.sac.trainer import SACTrainer
 from mlagents.trainers.sac.optimizer import SACOptimizer
-from mlagents.trainers.common.nn_policy import NNPolicy
+from mlagents.trainers.policy.nn_policy import NNPolicy
 from mlagents.trainers.agent_processor import AgentManagerQueue
 from mlagents.trainers.tests import mock_brain as mb
 from mlagents.trainers.tests.mock_brain import make_brain_parameters

ml-agents/mlagents/trainers/trainer/rl_trainer.py

 from typing import Dict
 from collections import defaultdict
 
-from mlagents.trainers.common.tf_optimizer import TFOptimizer
+from mlagents.trainers.optimizer.tf_optimizer import TFOptimizer
 from mlagents.trainers.buffer import AgentBuffer
 from mlagents.trainers.trainer import Trainer
 from mlagents.trainers.exception import UnityTrainerException

ml-agents/mlagents/trainers/trainer/trainer.py

 
 from mlagents_envs.timers import set_gauge
 from mlagents.model_serialization import export_policy_model, SerializationSettings
-from mlagents.trainers.tf_policy import TFPolicy
+from mlagents.trainers.policy.tf_policy import TFPolicy
 from mlagents.trainers.stats import StatsReporter
 from mlagents.trainers.trajectory import Trajectory
 from mlagents.trainers.agent_processor import AgentManagerQueue

ml-agents/mlagents/trainers/optimizer/tf_optimizer.py

 
 from mlagents.tf_utils.tf import tf
 from mlagents.trainers.buffer import AgentBuffer
-from mlagents.trainers.tf_policy import TFPolicy
-from mlagents.trainers.common.optimizer import Optimizer
+from mlagents.trainers.policy.tf_policy import TFPolicy
+from mlagents.trainers.optimizer import Optimizer
 from mlagents.trainers.trajectory import SplitObservations
 from mlagents.trainers.components.reward_signals.reward_signal_factory import (
     create_reward_signal,

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

 from mlagents.trainers.brain import BrainParameters
 from mlagents.trainers.models import EncoderType
 from mlagents.trainers.models import ModelUtils
-from mlagents.trainers.tf_policy import TFPolicy
-from mlagents.trainers.common.distributions import (
+from mlagents.trainers.policy.tf_policy import TFPolicy
+from mlagents.trainers.distributions import (
     GaussianDistribution,
     MultiCategoricalDistribution,
 )

ml-agents/mlagents/trainers/optimizer/__init__.py

+from mlagents.trainers.optimizer.optimizer import Optimizer  # noqa

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

+from mlagents.trainers.policy.policy import Policy  # noqa

ml-agents/mlagents/trainers/trainer/__init__.py

+from mlagents.trainers.trainer.trainer import Trainer  # noqa