Remove some dependencies of the trainers on UnityEnvironment (and use BaseEnv instead)

ml-agents/mlagents/trainers/learn.py

 ) -> Callable[[int, List[SideChannel]], BaseEnv]:
     def create_unity_environment(
         worker_id: int, side_channels: List[SideChannel]
-    ) -> UnityEnvironment:
+    ) -> BaseEnv:
         # Make sure that each environment gets a different seed
         env_seed = seed + worker_id
         return UnityEnvironment(

ml-agents/mlagents/trainers/subprocess_env_manager.py

 import enum
 import time
 
-from mlagents_envs.environment import UnityEnvironment
 from mlagents_envs.exception import (
     UnityCommunicationException,
     UnityTimeOutException,
     run_options: RunOptions,
     log_level: int = logging_util.INFO,
 ) -> None:
-    env_factory: Callable[
-        [int, List[SideChannel]], UnityEnvironment
-    ] = cloudpickle.loads(pickled_env_factory)
+    env_factory: Callable[[int, List[SideChannel]], BaseEnv] = cloudpickle.loads(
+        pickled_env_factory
+    )
     env_parameters = EnvironmentParametersChannel()
 
     engine_config = EngineConfig(
     training_analytics_channel: Optional[TrainingAnalyticsSideChannel] = None
     if worker_id == 0:
         training_analytics_channel = TrainingAnalyticsSideChannel()
-    env: UnityEnvironment = None
+    env: BaseEnv = None
     # Set log level. On some platforms, the logger isn't common with the
     # main process, so we need to set it again.
     logging_util.set_log_level(log_level)