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

Author: vincentpierre

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

@@ -35,14 +35,6 @@ class UnityActionException(UnityException):
     pass
 
 
-class SamplerException(UnityException):
-    """
-    Related to errors with the sampler actions.
-    """
-
-    pass
-
-
 class UnityTimeOutException(UnityException):
     """
     Related to errors with communication timeouts.

ml-agents-envs/mlagents/envs/action_info.py renamed to ml-agents/mlagents/trainers/action_info.py

@@ -6,8 +6,8 @@
 
 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/bc/trainer.py

@@ -1,4 +1,4 @@
-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-envs/mlagents/envs/brain.py renamed to ml-agents/mlagents/trainers/brain.py

@@ -6,7 +6,7 @@
 
 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-envs/mlagents/envs/brain_conversion_utils.py renamed to ml-agents/mlagents/trainers/brain_conversion_utils.py

@@ -2,7 +2,7 @@
 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/__init__.py

@@ -1,6 +1,6 @@
 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/curiosity/signal.py

@@ -3,7 +3,7 @@
 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/components/reward_signals/extrinsic/signal.py

@@ -5,7 +5,7 @@
 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/components/reward_signals/gail/signal.py

@@ -1,8 +1,8 @@
 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/demo_loader.py

@@ -41,3 +41,11 @@ class MetaCurriculumError(TrainerError):
     """
 
     pass
+
+
+class SamplerException(TrainerError):
+    """
+    Related to errors with the sampler actions.
+    """
+
+    pass

ml-agents-envs/mlagents/envs/env_manager.py renamed to ml-agents/mlagents/trainers/env_manager.py

@@ -18,10 +18,10 @@
 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.base_env import BaseEnv
-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
 
@@ -59,6 +59,7 @@ def from_argparse(args: Any) -> "CommandLineOptions":
 
 
 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/exception.py

@@ -6,7 +6,7 @@
 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/learn.py

@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
 
-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
 
 
 class Policy(ABC):

ml-agents/mlagents/trainers/models.py

@@ -4,7 +4,7 @@
 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-envs/mlagents/envs/policy.py renamed to ml-agents/mlagents/trainers/policy.py

@@ -5,7 +5,7 @@
 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/multi_gpu_policy.py

@@ -8,11 +8,11 @@
 
 import numpy as np
 
-from mlagents.envs.brain import BrainInfo
+from mlagents.trainers.brain import BrainInfo
 from mlagents.trainers.ppo.policy import PPOPolicy
 from mlagents.trainers.ppo.multi_gpu_policy import MultiGpuPPOPolicy, get_devices
 from mlagents.trainers.rl_trainer import RLTrainer, AllRewardsOutput
-from mlagents.envs.action_info import ActionInfoOutputs
+from mlagents.trainers.action_info import ActionInfoOutputs
 
 logger = logging.getLogger("mlagents.trainers")
 

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

@@ -3,8 +3,8 @@
 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/ppo/trainer.py

@@ -4,7 +4,7 @@
 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/rl_trainer.py

@@ -10,8 +10,8 @@
 
 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/sac/policy.py

@@ -2,7 +2,7 @@
 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/sac/trainer.py

@@ -1,12 +1,12 @@
 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.timers import timed
-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.side_channel.float_properties_channel import FloatPropertiesChannel
-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-envs/mlagents/envs/sampler_class.py renamed to ml-agents/mlagents/trainers/sampler_class.py

@@ -8,23 +8,23 @@
 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,
     hierarchical_timer,
     reset_timers,
     get_timer_root,
 )
-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,
     EngineConfig,
 )
 from mlagents.envs.side_channel.side_channel import SideChannel
-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-envs/mlagents/envs/simple_env_manager.py renamed to ml-agents/mlagents/trainers/simple_env_manager.py

@@ -1,7 +1,7 @@
 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-envs/mlagents/envs/subprocess_env_manager.py renamed to ml-agents/mlagents/trainers/subprocess_env_manager.py

@@ -14,7 +14,7 @@
 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/mock_brain.py

@@ -1,6 +1,6 @@
 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_bc.py

@@ -11,12 +11,12 @@
 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.environment import UnityEnvironment
 from mlagents.envs.mock_communicator import MockCommunicator
 from mlagents.trainers.tests import mock_brain as mb
 from mlagents.trainers.tests.mock_brain import make_brain_parameters
-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

@@ -1,12 +1,12 @@
 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 (
     UniformSampler,
     MultiRangeUniformSampler,
     GaussianSampler,
 )
-from mlagents.envs.exception import UnityException
+from mlagents.trainers.exception import TrainerError
 
 
 def sampler_config_1():
@@ -86,11 +86,11 @@ def incorrect_sampler_config():
 
 def test_incorrect_uniform_sampler():
     config = incorrect_uniform_sampler()
-    with pytest.raises(UnityException):
+    with pytest.raises(TrainerError):
         SamplerManager(config)
 
 
 def test_incorrect_sampler():
     config = incorrect_sampler_config()
-    with pytest.raises(UnityException):
+    with pytest.raises(TrainerError):
         SamplerManager(config)

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

@@ -15,9 +15,9 @@
     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-envs/mlagents/envs/tests/test_sampler_class.py renamed to ml-agents/mlagents/trainers/tests/test_sampler_class.py

@@ -3,7 +3,7 @@
 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/tests/test_simple_rl.py

@@ -6,8 +6,8 @@
 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-envs/mlagents/envs/tests/test_subprocess_env_manager.py renamed to ml-agents/mlagents/trainers/tests/test_subprocess_env_manager.py

@@ -113,7 +113,7 @@ def dummy_bad_config():
     )
 
 
-@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"
@@ -162,7 +162,7 @@ def mock_constructor(self, brain, trainer_parameters, training, load, seed, run_
         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"
@@ -221,7 +221,7 @@ def mock_constructor(
         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"