[feature] Fix TF tests, add --torch CLI option, allow run TF without torch installed (#4305)

Author: Ervin T

.circleci/config.yml

@@ -70,7 +70,7 @@ jobs:
             . venv/bin/activate
             mkdir test-reports
             pip freeze > test-reports/pip_versions.txt
-            pytest -n 2 --cov=ml-agents --cov=ml-agents-envs --cov=gym-unity --cov-report html --junitxml=test-reports/junit.xml -p no:warnings
+            pytest --cov=ml-agents --cov=ml-agents-envs --cov=gym-unity --cov-report html --junitxml=test-reports/junit.xml -p no:warnings
 
       - run:
           name: Verify there are no hidden/missing metafiles.

experiment_torch.py

@@ -98,8 +98,8 @@ def run_experiment(
         evaluate_count = evaluate["TorchPolicy.evaluate"]["count"]
     else:
         if algo == "ppo":
-            update_total = update["TFPPOOptimizer.update"]["total"]
-            update_count = update["TFPPOOptimizer.update"]["count"]
+            update_total = update["PPOOptimizer.update"]["total"]
+            update_count = update["PPOOptimizer.update"]["count"]
         else:
             update_total = update["SACTrainer._update_policy"]["total"]
             update_count = update["SACTrainer._update_policy"]["count"]

ml-agents/mlagents/trainers/cli_utils.py

@@ -168,6 +168,13 @@ def _create_parser() -> argparse.ArgumentParser:
         action=DetectDefaultStoreTrue,
         help="Forces training using CPU only",
     )
+    argparser.add_argument(
+        "--torch",
+        default=False,
+        action=DetectDefaultStoreTrue,
+        help="(Experimental) Use the PyTorch framework instead of TensorFlow. Install PyTorch "
+        "before using this option",
+    )
 
     eng_conf = argparser.add_argument_group(title="Engine Configuration")
     eng_conf.add_argument(

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

@@ -9,7 +9,7 @@
 from mlagents.trainers.settings import TrainerSettings, PPOSettings
 
 
-class TFPPOOptimizer(TFOptimizer):
+class PPOOptimizer(TFOptimizer):
     def __init__(self, policy: TFPolicy, trainer_params: TrainerSettings):
         """
         Takes a Policy and a Dict of trainer parameters and creates an Optimizer around the policy.

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

@@ -10,19 +10,25 @@
 from mlagents_envs.logging_util import get_logger
 from mlagents_envs.base_env import BehaviorSpec
 from mlagents.trainers.trainer.rl_trainer import RLTrainer
-from mlagents.trainers.policy.torch_policy import TorchPolicy
 from mlagents.trainers.policy import Policy
 from mlagents.trainers.policy.tf_policy import TFPolicy
-from mlagents.trainers.ppo.optimizer_torch import TorchPPOOptimizer
-from mlagents.trainers.ppo.optimizer_tf import TFPPOOptimizer
+from mlagents.trainers.ppo.optimizer_tf import PPOOptimizer
 from mlagents.trainers.trajectory import Trajectory
 from mlagents.trainers.behavior_id_utils import BehaviorIdentifiers
 from mlagents.trainers.settings import (
     TrainerSettings,
     PPOSettings,
     TestingConfiguration,
+    FrameworkType,
 )
 
+try:
+    from mlagents.trainers.policy.torch_policy import TorchPolicy
+    from mlagents.trainers.ppo.optimizer_torch import TorchPPOOptimizer
+except ModuleNotFoundError:
+    TorchPolicy = None  # type: ignore
+    TorchPPOOptimizer = None  # type: ignore
+
 
 logger = get_logger(__name__)
 
@@ -58,7 +64,6 @@ def __init__(
         )
         self.load = load
         self.seed = seed
-        self.framework = "torch" if TestingConfiguration.use_torch else "tf"
         if TestingConfiguration.max_steps > 0:
             self.trainer_settings.max_steps = TestingConfiguration.max_steps
         self.policy: Policy = None  # type: ignore
@@ -254,12 +259,12 @@ def add_policy(
             )
         self.policy = policy
         self.policies[parsed_behavior_id.behavior_id] = policy
-        if self.framework == "torch":
+        if self.framework == FrameworkType.PYTORCH:
             self.optimizer = TorchPPOOptimizer(  # type: ignore
                 self.policy, self.trainer_settings  # type: ignore
             )  # type: ignore
         else:
-            self.optimizer = TFPPOOptimizer(  # type: ignore
+            self.optimizer = PPOOptimizer(  # type: ignore
                 self.policy, self.trainer_settings  # type: ignore
             )  # type: ignore
         for _reward_signal in self.optimizer.reward_signals.keys():

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

@@ -18,10 +18,14 @@
 from mlagents.trainers.trainer.rl_trainer import RLTrainer
 from mlagents.trainers.trajectory import Trajectory, SplitObservations
 from mlagents.trainers.behavior_id_utils import BehaviorIdentifiers
-from mlagents.trainers.policy.torch_policy import TorchPolicy
-from mlagents.trainers.sac.optimizer_torch import TorchSACOptimizer
-from mlagents.trainers.settings import TrainerSettings, SACSettings
+from mlagents.trainers.settings import TrainerSettings, SACSettings, FrameworkType
 
+try:
+    from mlagents.trainers.policy.torch_policy import TorchPolicy
+    from mlagents.trainers.sac.optimizer_torch import TorchSACOptimizer
+except ModuleNotFoundError:
+    TorchPolicy = None  # type: ignore
+    TorchSACOptimizer = None  # type: ignore
 
 logger = get_logger(__name__)
 
@@ -353,7 +357,7 @@ def add_policy(
             )
         self.policy = policy
         self.policies[parsed_behavior_id.behavior_id] = policy
-        if self.framework == "torch":
+        if self.framework == FrameworkType.PYTORCH:
             self.optimizer = TorchSACOptimizer(  # type: ignore
                 self.policy, self.trainer_settings  # type: ignore
             )  # type: ignore

ml-agents/mlagents/trainers/settings.py

@@ -524,6 +524,11 @@ def to_settings(self) -> type:
         return _mapping[self]
 
 
+class FrameworkType(Enum):
+    TENSORFLOW: str = "tensorflow"
+    PYTORCH: str = "pytorch"
+
+
 @attr.s(auto_attribs=True)
 class TrainerSettings(ExportableSettings):
     trainer_type: TrainerType = TrainerType.PPO
@@ -546,6 +551,7 @@ def _set_default_hyperparameters(self):
     threaded: bool = True
     self_play: Optional[SelfPlaySettings] = None
     behavioral_cloning: Optional[BehavioralCloningSettings] = None
+    framework: FrameworkType = FrameworkType.TENSORFLOW
 
     cattr.register_structure_hook(
         Dict[RewardSignalType, RewardSignalSettings], RewardSignalSettings.structure
@@ -713,7 +719,13 @@ def from_argparse(args: argparse.Namespace) -> "RunOptions":
                     configured_dict["engine_settings"][key] = val
                 else:  # Base options
                     configured_dict[key] = val
-        return RunOptions.from_dict(configured_dict)
+
+        # Apply --torch retroactively
+        final_runoptions = RunOptions.from_dict(configured_dict)
+        if "torch" in DetectDefault.non_default_args:
+            for trainer_set in final_runoptions.behaviors.values():
+                trainer_set.framework = FrameworkType.PYTORCH
+        return final_runoptions
 
     @staticmethod
     def from_dict(options_dict: Dict[str, Any]) -> "RunOptions":

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

@@ -8,7 +8,7 @@
 from mlagents.trainers.behavior_id_utils import BehaviorIdentifiers
 
 from mlagents.trainers.ppo.trainer import PPOTrainer, discount_rewards
-from mlagents.trainers.ppo.optimizer_tf import TFPPOOptimizer
+from mlagents.trainers.ppo.optimizer_tf import PPOOptimizer
 from mlagents.trainers.policy.tf_policy import TFPolicy
 from mlagents.trainers.agent_processor import AgentManagerQueue
 from mlagents.trainers.tests import mock_brain as mb
@@ -52,7 +52,7 @@ def _create_ppo_optimizer_ops_mock(dummy_config, use_rnn, use_discrete, use_visu
     policy = TFPolicy(
         0, mock_specs, trainer_settings, "test", False, create_tf_graph=False
     )
-    optimizer = TFPPOOptimizer(policy, trainer_settings)
+    optimizer = PPOOptimizer(policy, trainer_settings)
     return optimizer
 
 

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

@@ -4,7 +4,7 @@
 import mlagents.trainers.tests.mock_brain as mb
 from mlagents.trainers.policy.tf_policy import TFPolicy
 from mlagents.trainers.sac.optimizer import SACOptimizer
-from mlagents.trainers.ppo.optimizer_tf import TFPPOOptimizer
+from mlagents.trainers.ppo.optimizer_tf import PPOOptimizer
 from mlagents.trainers.tests.test_simple_rl import PPO_CONFIG, SAC_CONFIG
 from mlagents.trainers.settings import (
     GAILSettings,
@@ -75,7 +75,7 @@ def create_optimizer_mock(
     if trainer_settings.trainer_type == TrainerType.SAC:
         optimizer = SACOptimizer(policy, trainer_settings)
     else:
-        optimizer = TFPPOOptimizer(policy, trainer_settings)
+        optimizer = PPOOptimizer(policy, trainer_settings)
     return optimizer
 
 

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

@@ -25,7 +25,10 @@ def _update_policy(self):
     def add_policy(self, mock_behavior_id, mock_policy):
         self.policies[mock_behavior_id] = mock_policy
 
-    def create_policy(self):
+    def create_tf_policy(self):
+        return mock.Mock()
+
+    def create_torch_policy(self):
         return mock.Mock()
 
     def _process_trajectory(self, trajectory):