Moving Env Manager to Trainers (#3062) The Env Manager is only used by the trainer codebase. The entry point to interact with an environment is UnityEnvironment.

* Moving Env Manager to Trainers

* fix pylint madness

ml-agents-envs/mlagents/envs/exception.py

     pass
 
 
-class SamplerException(UnityException):
-    """
-    Related to errors with the sampler actions.
-    """
-
-    pass
-
-
 class UnityTimeOutException(UnityException):
     """
     Related to errors with communication timeouts.

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

 
 import numpy as np
 
-from mlagents.envs.brain import BrainInfo
-from mlagents.envs.action_info import ActionInfoOutputs
+from mlagents.trainers.brain import BrainInfo
+from mlagents.trainers.action_info import ActionInfoOutputs
 from mlagents.trainers.bc.policy import BCPolicy
 from mlagents.trainers.buffer import AgentBuffer
 from mlagents.trainers.agent_processor import ProcessingBuffer

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

 
 from mlagents.tf_utils import tf
 
-from mlagents.envs.brain import BrainInfo
+from mlagents.trainers.brain import BrainInfo
 from mlagents.trainers.trainer import UnityTrainerException
 from mlagents.trainers.tf_policy import TFPolicy
 from mlagents.trainers.models import LearningModel

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

 import numpy as np
 from mlagents.tf_utils import tf
 
-from mlagents.envs.brain import BrainInfo
+from mlagents.trainers.brain import BrainInfo
 
 from mlagents.trainers.components.reward_signals import RewardSignal, RewardSignalResult
 from mlagents.trainers.components.reward_signals.curiosity.model import CuriosityModel

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

 from typing import Any, Dict, List
 import numpy as np
-from mlagents.envs.brain import BrainInfo
+from mlagents.trainers.brain import BrainInfo
 
 from mlagents.trainers.components.reward_signals import RewardSignal, RewardSignalResult
 

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

 import numpy as np
 from mlagents.tf_utils import tf
 
-from mlagents.envs.brain import BrainInfo
+from mlagents.trainers.brain import BrainInfo
 from mlagents.trainers.components.reward_signals import RewardSignal, RewardSignalResult
 from mlagents.trainers.tf_policy import TFPolicy
 from mlagents.trainers.models import LearningModel

ml-agents/mlagents/trainers/demo_loader.py

 import numpy as np
 from mlagents.trainers.buffer import AgentBuffer
 from mlagents.trainers.agent_processor import ProcessingBuffer
-from mlagents.envs.brain import BrainParameters, BrainInfo
+from mlagents.trainers.brain import BrainParameters, BrainInfo
 from mlagents.envs.communicator_objects.agent_info_action_pair_pb2 import (
     AgentInfoActionPairProto,
 )

ml-agents/mlagents/trainers/exception.py

     """
 
     pass
+
+
+class SamplerException(TrainerError):
+    """
+    Related to errors with the sampler actions.
+    """
+
+    pass

ml-agents/mlagents/trainers/learn.py

 from mlagents.trainers.meta_curriculum import MetaCurriculum
 from mlagents.trainers.trainer_util import load_config, TrainerFactory
 from mlagents.envs.environment import UnityEnvironment
-from mlagents.envs.sampler_class import SamplerManager
-from mlagents.envs.exception import SamplerException
+from mlagents.trainers.sampler_class import SamplerManager
+from mlagents.trainers.exception import SamplerException
-from mlagents.envs.subprocess_env_manager import SubprocessEnvManager
+from mlagents.trainers.subprocess_env_manager import SubprocessEnvManager
 from mlagents.envs.side_channel.side_channel import SideChannel
 from mlagents.envs.side_channel.engine_configuration_channel import EngineConfig
 
 
 
 def get_version_string() -> str:
+    # pylint: disable=no-member
     return f""" Version information:
   ml-agents: {mlagents.trainers.__version__},
   ml-agents-envs: {mlagents.envs.__version__},

ml-agents/mlagents/trainers/models.py

 from mlagents.tf_utils import tf
 
 from mlagents.trainers.trainer import UnityTrainerException
-from mlagents.envs.brain import CameraResolution
+from mlagents.trainers.brain import CameraResolution
 
 logger = logging.getLogger("mlagents.trainers")
 

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

 from mlagents.tf_utils import tf
 
 from tensorflow.python.client import device_lib
-from mlagents.envs.brain import BrainParameters
+from mlagents.trainers.brain import BrainParameters
 from mlagents.envs.timers import timed
 from mlagents.trainers.models import EncoderType, LearningRateSchedule
 from mlagents.trainers.ppo.policy import PPOPolicy

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

 from mlagents.tf_utils import tf
 
 from mlagents.envs.timers import timed
-from mlagents.envs.brain import BrainInfo, BrainParameters
+from mlagents.trainers.brain import BrainInfo, BrainParameters
 from mlagents.trainers.models import EncoderType, LearningRateSchedule
 from mlagents.trainers.ppo.models import PPOModel
 from mlagents.trainers.tf_policy import TFPolicy

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

 
 import numpy as np
 
-from mlagents.envs.brain import BrainInfo
+from mlagents.trainers.brain import BrainInfo
-from mlagents.envs.action_info import ActionInfoOutputs
+from mlagents.trainers.action_info import ActionInfoOutputs
 
 logger = logging.getLogger("mlagents.trainers")
 

ml-agents/mlagents/trainers/rl_trainer.py

 from typing import Dict, List, Any, NamedTuple
 import numpy as np
 
-from mlagents.envs.brain import BrainInfo
-from mlagents.envs.action_info import ActionInfoOutputs
+from mlagents.trainers.brain import BrainInfo
+from mlagents.trainers.action_info import ActionInfoOutputs
 from mlagents.trainers.buffer import AgentBuffer
 from mlagents.trainers.agent_processor import ProcessingBuffer
 from mlagents.trainers.trainer import Trainer, UnityTrainerException

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

 from mlagents.tf_utils import tf
 
 from mlagents.envs.timers import timed
-from mlagents.envs.brain import BrainInfo, BrainParameters
+from mlagents.trainers.brain import BrainInfo, BrainParameters
 from mlagents.trainers.models import EncoderType, LearningRateSchedule
 from mlagents.trainers.sac.models import SACModel
 from mlagents.trainers.tf_policy import TFPolicy

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

 
 import numpy as np
 
-from mlagents.envs.brain import BrainInfo
-from mlagents.envs.action_info import ActionInfoOutputs
+from mlagents.trainers.brain import BrainInfo
+from mlagents.trainers.action_info import ActionInfoOutputs
 from mlagents.envs.timers import timed
 from mlagents.trainers.sac.policy import SACPolicy
 from mlagents.trainers.rl_trainer import RLTrainer, AllRewardsOutput

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

 import unittest.mock as mock
 import numpy as np
 
-from mlagents.envs.brain import CameraResolution, BrainParameters
+from mlagents.trainers.brain import CameraResolution, BrainParameters
 from mlagents.trainers.buffer import AgentBuffer
 from mlagents.trainers.agent_processor import ProcessingBuffer
 

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

 from mlagents.envs.mock_communicator import MockCommunicator
 from mlagents.trainers.tests.mock_brain import make_brain_parameters
 from mlagents.envs.environment import UnityEnvironment
-from mlagents.envs.brain_conversion_utils import (
+from mlagents.trainers.brain_conversion_utils import (
     step_result_to_brain_info,
     group_spec_to_brain_parameters,
 )

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

 from mlagents.trainers.tf_policy import TFPolicy
-from mlagents.envs.brain import BrainInfo
-from mlagents.envs.action_info import ActionInfo
+from mlagents.trainers.brain import BrainInfo
+from mlagents.trainers.action_info import ActionInfo
 from unittest.mock import MagicMock
 import numpy as np
 

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

 from mlagents.trainers.ppo.policy import PPOPolicy
 from mlagents.trainers.rl_trainer import AllRewardsOutput
 from mlagents.trainers.components.reward_signals import RewardSignalResult
-from mlagents.envs.brain import BrainParameters
+from mlagents.trainers.brain import BrainParameters
-from mlagents.envs.brain_conversion_utils import (
+from mlagents.trainers.brain_conversion_utils import (
     step_result_to_brain_info,
     group_spec_to_brain_parameters,
 )

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

     BatchedStepResult,
     ActionType,
 )
-from mlagents.envs.brain import BrainParameters
-from mlagents.envs.simple_env_manager import SimpleEnvManager
-from mlagents.envs.sampler_class import SamplerManager
+from mlagents.trainers.brain import BrainParameters
+from mlagents.trainers.simple_env_manager import SimpleEnvManager
+from mlagents.trainers.sampler_class import SamplerManager
 from mlagents.envs.side_channel.float_properties_channel import FloatPropertiesChannel
 
 BRAIN_NAME = __name__

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

 import pytest
 
 from mlagents.trainers.trainer_controller import TrainerController
-from mlagents.envs.subprocess_env_manager import EnvironmentStep
-from mlagents.envs.sampler_class import SamplerManager
+from mlagents.trainers.subprocess_env_manager import EnvironmentStep
+from mlagents.trainers.sampler_class import SamplerManager
 
 
 @pytest.fixture

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

     )
 
 
-@patch("mlagents.envs.brain.BrainParameters")
+@patch("mlagents.trainers.brain.BrainParameters")
 def test_initialize_trainer_parameters_override_defaults(BrainParametersMock):
     summaries_dir = "test_dir"
     run_id = "testrun"
         assert isinstance(trainers["testbrain"], OfflineBCTrainer)
 
 
-@patch("mlagents.envs.brain.BrainParameters")
+@patch("mlagents.trainers.brain.BrainParameters")
 def test_initialize_ppo_trainer(BrainParametersMock):
     brain_params_mock = BrainParametersMock()
     BrainParametersMock.return_value.brain_name = "testbrain"
         assert isinstance(trainers["testbrain"], PPOTrainer)
 
 
-@patch("mlagents.envs.brain.BrainParameters")
+@patch("mlagents.trainers.brain.BrainParameters")
 def test_initialize_invalid_trainer_raises_exception(BrainParametersMock):
     summaries_dir = "test_dir"
     run_id = "testrun"

ml-agents/mlagents/trainers/tf_policy.py

 from mlagents.tf_utils import tf
 
 from mlagents.envs.exception import UnityException
-from mlagents.envs.policy import Policy
-from mlagents.envs.action_info import ActionInfo
+from mlagents.trainers.policy import Policy
+from mlagents.trainers.action_info import ActionInfo
-from mlagents.envs.brain import BrainInfo
+from mlagents.trainers.brain import BrainInfo
 
 
 logger = logging.getLogger("mlagents.trainers")

ml-agents/mlagents/trainers/trainer.py

 import numpy as np
 from collections import deque, defaultdict
 
-from mlagents.envs.action_info import ActionInfoOutputs
+from mlagents.trainers.action_info import ActionInfoOutputs
-from mlagents.envs.brain import BrainParameters, BrainInfo
+from mlagents.trainers.brain import BrainParameters, BrainInfo
 
 LOGGER = logging.getLogger("mlagents.trainers")
 

ml-agents/mlagents/trainers/trainer_controller.py

 from mlagents.tf_utils import tf
 from time import time
 
-from mlagents.envs.env_manager import EnvironmentStep
-from mlagents.envs.env_manager import EnvManager
+from mlagents.trainers.env_manager import EnvManager, EnvironmentStep
-from mlagents.envs.sampler_class import SamplerManager
+from mlagents.trainers.sampler_class import SamplerManager
 from mlagents.envs.timers import hierarchical_timer, get_timer_tree, timed
 from mlagents.trainers.trainer import Trainer, TrainerMetrics
 from mlagents.trainers.meta_curriculum import MetaCurriculum

ml-agents/mlagents/trainers/trainer_util.py

 from mlagents.trainers.meta_curriculum import MetaCurriculum
 from mlagents.envs.exception import UnityEnvironmentException
 from mlagents.trainers.trainer import Trainer
-from mlagents.envs.brain import BrainParameters
+from mlagents.trainers.brain import BrainParameters
 from mlagents.trainers.ppo.trainer import PPOTrainer
 from mlagents.trainers.sac.trainer import SACTrainer
 from mlagents.trainers.bc.offline_trainer import OfflineBCTrainer

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

 import pytest
 
-from mlagents.envs.sampler_class import SamplerManager
-from mlagents.envs.sampler_class import (
+from mlagents.trainers.sampler_class import SamplerManager
+from mlagents.trainers.sampler_class import (
-from mlagents.envs.exception import UnityException
+from mlagents.trainers.exception import TrainerError
 
 
 def sampler_config_1():
 
 def test_incorrect_uniform_sampler():
     config = incorrect_uniform_sampler()
-    with pytest.raises(UnityException):
+    with pytest.raises(TrainerError):
-    with pytest.raises(UnityException):
+    with pytest.raises(TrainerError):
         SamplerManager(config)

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

 import unittest
 from queue import Empty as EmptyQueue
 
-from mlagents.envs.subprocess_env_manager import (
+from mlagents.trainers.subprocess_env_manager import (
     SubprocessEnvManager,
     EnvironmentResponse,
     StepResponse,

ml-agents/mlagents/trainers/sampler_class.py

 from typing import Union, Optional, Type, List, Dict, Any
 from abc import ABC, abstractmethod
 
-from .exception import SamplerException
+from mlagents.trainers.exception import SamplerException
 
 
 class Sampler(ABC):

ml-agents/mlagents/trainers/env_manager.py

 from abc import ABC, abstractmethod
 from typing import List, Dict, NamedTuple, Optional
-from mlagents.envs.brain import AllBrainInfo, BrainParameters
-from mlagents.envs.policy import Policy
-from mlagents.envs.action_info import ActionInfo
+from mlagents.trainers.brain import AllBrainInfo, BrainParameters
+from mlagents.trainers.policy import Policy
+from mlagents.trainers.action_info import ActionInfo
 
 
 class EnvironmentStep(NamedTuple):

ml-agents/mlagents/trainers/simple_env_manager.py

 from typing import Dict, List
 
 from mlagents.envs.base_env import BaseEnv
-from mlagents.envs.env_manager import EnvManager, EnvironmentStep
+from mlagents.trainers.env_manager import EnvManager, EnvironmentStep
-from mlagents.envs.action_info import ActionInfo
-from mlagents.envs.brain import BrainParameters, AllBrainInfo
+from mlagents.trainers.action_info import ActionInfo
+from mlagents.trainers.brain import BrainParameters, AllBrainInfo
-from mlagents.envs.brain_conversion_utils import (
+from mlagents.trainers.brain_conversion_utils import (
     step_result_to_brain_info,
     group_spec_to_brain_parameters,
 )

ml-agents/mlagents/trainers/subprocess_env_manager.py

 from multiprocessing.connection import Connection
 from queue import Empty as EmptyQueueException
 from mlagents.envs.base_env import BaseEnv
-from mlagents.envs.env_manager import EnvManager, EnvironmentStep
+from mlagents.trainers.env_manager import EnvManager, EnvironmentStep
 from mlagents.envs.timers import (
     TimerNode,
     timed,
 )
-from mlagents.envs.brain import AllBrainInfo, BrainParameters
-from mlagents.envs.action_info import ActionInfo
+from mlagents.trainers.brain import AllBrainInfo, BrainParameters
+from mlagents.trainers.action_info import ActionInfo
 from mlagents.envs.side_channel.float_properties_channel import FloatPropertiesChannel
 from mlagents.envs.side_channel.engine_configuration_channel import (
     EngineConfigurationChannel,
-from mlagents.envs.brain_conversion_utils import (
+from mlagents.trainers.brain_conversion_utils import (
     step_result_to_brain_info,
     group_spec_to_brain_parameters,
 )

ml-agents/mlagents/trainers/brain_conversion_utils.py

-from mlagents.envs.brain import BrainInfo, BrainParameters, CameraResolution
+from mlagents.trainers.brain import BrainInfo, BrainParameters, CameraResolution
 from mlagents.envs.base_env import BatchedStepResult, AgentGroupSpec
 from mlagents.envs.exception import UnityEnvironmentException
 import numpy as np

ml-agents/mlagents/trainers/policy.py

+from abc import ABC, abstractmethod
+
+from mlagents.trainers.brain import BrainInfo
+from mlagents.trainers.action_info import ActionInfo
+
+
+class Policy(ABC):
+    @abstractmethod
+    def get_action(self, brain_info: BrainInfo) -> ActionInfo:
+        pass

ml-agents-envs/mlagents/envs/policy.py

-from abc import ABC, abstractmethod
-
-from mlagents.envs.brain import BrainInfo
-from mlagents.envs.action_info import ActionInfo
-
-
-class Policy(ABC):
-    @abstractmethod
-    def get_action(self, brain_info: BrainInfo) -> ActionInfo:
-        pass