[RLlib] Assert correct policy class being used in Worker. (#7769)

rllib/contrib/alpha_zero/core/alpha_zero_trainer.py

@@ -96,7 +96,9 @@ DEFAULT_CONFIG = with_common_config({
     # === Callbacks ===
     "callbacks": {
         "on_episode_start": on_episode_start,
-    }
+    },
+
+    "use_pytorch": True,
 })
 
 

rllib/evaluation/worker_set.py

@@ -5,6 +5,7 @@ import ray
 from ray.rllib.utils.annotations import DeveloperAPI
 from ray.rllib.evaluation.rollout_worker import RolloutWorker, \
     _validate_multiagent_config
+from ray.rllib.policy import Policy, TorchPolicy
 from ray.rllib.offline import NoopOutput, JsonReader, MixedInput, JsonWriter, \
     ShuffledInput
 from ray.rllib.utils import merge_dicts, try_import_tf
@@ -233,7 +234,7 @@ class WorkerSet:
                     tmp[k] = (policy, v[1], v[2], v[3])
             policy = tmp
 
-        return cls(
+        worker = cls(
             env_creator,
             policy,
             policy_mapping_fn=config["multiagent"]["policy_mapping_fn"],
@@ -269,3 +270,26 @@ class WorkerSet:
             seed=(config["seed"] + worker_index)
             if config["seed"] is not None else None,
             _fake_sampler=config.get("_fake_sampler", False))
+
+        # Check for correct policy class (only locally, remote Workers should
+        # create the exact same Policy types).
+        if type(worker) is RolloutWorker:
+            actual_class = type(worker.get_policy())
+
+            # Pytorch case: Policy must be a TorchPolicy.
+            if config["use_pytorch"]:
+                assert issubclass(actual_class, TorchPolicy), \
+                    "Worker policy must be subclass of `TorchPolicy`, " \
+                    "but is {}!".format(actual_class.__name__)
+            # non-Pytorch case:
+            # Policy may be None AND must not be a TorchPolicy.
+            else:
+                assert issubclass(actual_class, type(None)) or \
+                       (issubclass(actual_class, Policy) and
+                        not issubclass(actual_class, TorchPolicy)), "Worker " \
+                       "policy must be subclass of `Policy`, but NOT " \
+                       "`TorchPolicy` (your class={})! If you have a torch " \
+                       "Trainer, make sure to set `use_pytorch=True` in " \
+                       "your Trainer's config)!".format(actual_class.__name__)
+
+        return worker

rllib/examples/custom_torch_policy.py

@@ -37,4 +37,5 @@ if __name__ == "__main__":
         config={
             "env": "CartPole-v0",
             "num_workers": 2,
+            "use_pytorch": True
         })

rllib/utils/exploration/tests/test_explorations.py

@@ -30,7 +30,7 @@ def do_test_explorations(run,
         config["num_workers"] = 0
 
     # Test all frameworks.
-    for fw in ["torch", "eager", "tf"]:
+    for fw in ["tf", "eager", "torch"]:
         if fw == "torch" and \
                 run in [ddpg.DDPGTrainer, dqn.DQNTrainer, dqn.SimpleQTrainer,
                         impala.ImpalaTrainer, sac.SACTrainer, td3.TD3Trainer]:
@@ -39,8 +39,8 @@ def do_test_explorations(run,
             continue
 
         print("Testing {} in framework={}".format(run, fw))
-        config["eager"] = (fw == "eager")
-        config["use_pytorch"] = (fw == "torch")
+        config["eager"] = fw == "eager"
+        config["use_pytorch"] = fw == "torch"
 
         # Test for both the default Agent's exploration AND the `Random`
         # exploration class.
@@ -52,9 +52,10 @@ def do_test_explorations(run,
                 config["exploration_config"] = {"type": "Random"}
             print("exploration={}".format(exploration or "default"))
 
-            eager_mode_ctx = eager_mode()
+            eager_ctx = None
             if fw == "eager":
-                eager_mode_ctx.__enter__()
+                eager_ctx = eager_mode()
+                eager_ctx.__enter__()
                 assert tf.executing_eagerly()
             elif fw == "tf":
                 assert not tf.executing_eagerly()
@@ -64,7 +65,7 @@ def do_test_explorations(run,
             # Make sure all actions drawn are the same, given same
             # observations.
             actions = []
-            for _ in range(100):
+            for _ in range(50):
                 actions.append(
                     trainer.compute_action(
                         observation=dummy_obs,
@@ -91,8 +92,8 @@ def do_test_explorations(run,
             # Check that the stddev is not 0.0 (values differ).
             check(np.std(actions), 0.0, false=True)
 
-            if fw == "eager":
-                eager_mode_ctx.__exit__(None, None, None)
+            if eager_ctx:
+                eager_ctx.__exit__(None, None, None)
 
 
 class TestExplorations(unittest.TestCase):