[RLlib] Add torch flag to train.py (#6807)

doc/source/rllib-concepts.rst

@@ -418,11 +418,22 @@ Finally, note that you do not have to use ``build_tf_policy`` to define a Tensor
 Building Policies in TensorFlow Eager
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-Policies built with ``build_tf_policy`` (most of the reference algorithms are) can be run in eager mode by setting the ``"eager": True`` / ``"eager_tracing": True`` config options or using ``rllib train --eager [--trace]``. This will tell RLlib to execute the model forward pass, action distribution, loss, and stats functions in eager mode.
+Policies built with ``build_tf_policy`` (most of the reference algorithms are)
+can be run in eager mode by setting
+the ``"eager": True`` / ``"eager_tracing": True`` config options or
+using ``rllib train --eager [--trace]``.
+This will tell RLlib to execute the model forward pass, action distribution,
+loss, and stats functions in eager mode.
 
-Eager mode makes debugging much easier, since you can now use normal Python functions such as ``print()`` to inspect intermediate tensor values. However, it can be slower than graph mode unless tracing is enabled.
+Eager mode makes debugging much easier, since you can now use line-by-line
+debugging with breakpoints or Python ``print()`` to inspect
+intermediate tensor values.
+However, eager can be slower than graph mode unless tracing is enabled.
 
-You can also selectively leverage eager operations within graph mode execution with `tf.py_function <https://www.tensorflow.org/api_docs/python/tf/py_function>`__. Here's an example of using eager ops embedded `within a loss function <https://github.com/ray-project/ray/blob/master/rllib/examples/eager_execution.py>`__.
+You can also selectively leverage eager operations within graph mode
+execution with `tf.py_function <https://www.tensorflow.org/api_docs/python/tf/py_function>`__.
+Here's an example of using eager ops embedded
+`within a loss function <https://github.com/ray-project/ray/blob/master/rllib/examples/eager_execution.py>`__.
 
 Building Policies in PyTorch
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~

doc/source/rllib-training.rst

@@ -618,9 +618,26 @@ The ``"monitor": true`` config can be used to save Gym episode videos to the res
 Eager Mode
 ~~~~~~~~~~
 
-Policies built with ``build_tf_policy`` (most of the reference algorithms are) can be run in eager mode by setting the ``"eager": True`` / ``"eager_tracing": True`` config options or using ``rllib train --eager [--trace]``. This will tell RLlib to execute the model forward pass, action distribution, loss, and stats functions in eager mode.
+Policies built with ``build_tf_policy`` (most of the reference algorithms are)
+can be run in eager mode by setting the
+``"eager": True`` / ``"eager_tracing": True`` config options or using
+``rllib train --eager [--trace]``.
+This will tell RLlib to execute the model forward pass, action distribution,
+loss, and stats functions in eager mode.
+
+Eager mode makes debugging much easier, since you can now use line-by-line
+debugging with breakpoints or Python ``print()`` to inspect
+intermediate tensor values.
+However, eager can be slower than graph mode unless tracing is enabled.
+
+Using PyTorch
+~~~~~~~~~~~~~
+
+Trainers that have an implemented TorchPolicy, will allow you to run
+`rllib train` using the the command line ``--torch`` flag.
+Algorithms that do not have a torch version yet will complain with an error in
+this case.
 
-Eager mode makes debugging much easier, since you can now use normal Python functions such as ``print()`` to inspect intermediate tensor values. However, it can be slower than graph mode unless tracing is enabled.
 
 Episode Traces
 ~~~~~~~~~~~~~~

rllib/agents/a3c/a3c.py

@@ -8,8 +8,6 @@ from ray.rllib.optimizers import AsyncGradientsOptimizer
 DEFAULT_CONFIG = with_common_config({
     # Size of rollout batch
     "sample_batch_size": 10,
-    # Use PyTorch as framework - no LSTM support
-    "use_pytorch": False,
     # GAE(gamma) parameter
     "lambda": 1.0,
     # Max global norm for each gradient calculated by worker

rllib/agents/ars/ars.py

@@ -162,6 +162,12 @@ class ARSTrainer(Trainer):
 
     @override(Trainer)
     def _init(self, config, env_creator):
+        # PyTorch check.
+        if config["use_pytorch"]:
+            raise ValueError(
+                "ARS does not support PyTorch yet! Use tf instead."
+            )
+
         env = env_creator(config["env_config"])
         from ray.rllib import models
         preprocessor = models.ModelCatalog.get_preprocessor(env)

rllib/agents/ddpg/ddpg.py

@@ -210,9 +210,16 @@ def add_pure_exploration_phase(trainer):
     update_worker_explorations(trainer)
 
 
+def validate_config(config):
+    # PyTorch check.
+    if config["use_pytorch"]:
+        raise ValueError("DDPG does not support PyTorch yet! Use tf instead.")
+
+
 DDPGTrainer = GenericOffPolicyTrainer.with_updates(
     name="DDPG",
     default_config=DEFAULT_CONFIG,
     default_policy=DDPGTFPolicy,
+    validate_config=validate_config,
     before_init=setup_ddpg_exploration,
     before_train_step=add_pure_exploration_phase)

rllib/agents/dqn/dqn.py

@@ -142,6 +142,10 @@ def check_config_and_setup_param_noise(config):
     adds the necessary callbacks to support parameter space noise exploration.
     """
 
+    # PyTorch check.
+    if config["use_pytorch"]:
+        raise ValueError("DQN does not support PyTorch yet! Use tf instead.")
+
     # Update effective batch size to include n-step
     adjusted_batch_size = max(config["sample_batch_size"],
                               config.get("n_step", 1))

rllib/agents/es/es.py

@@ -168,6 +168,12 @@ class ESTrainer(Trainer):
 
     @override(Trainer)
     def _init(self, config, env_creator):
+        # PyTorch check.
+        if config["use_pytorch"]:
+            raise ValueError(
+                "ES does not support PyTorch yet! Use tf instead."
+            )
+
         policy_params = {"action_noise_std": 0.01}
 
         env = env_creator(config["env_config"])

rllib/agents/impala/impala.py

@@ -92,6 +92,11 @@ def choose_policy(config):
 
 
 def validate_config(config):
+    # PyTorch check.
+    if config["use_pytorch"]:
+        raise ValueError(
+            "IMPALA does not support PyTorch yet! Use tf instead."
+        )
     if config["entropy_coeff"] < 0:
         raise DeprecationWarning("entropy_coeff must be >= 0")
 

rllib/agents/marwil/marwil.py

@@ -44,8 +44,16 @@ def make_optimizer(workers, config):
     )
 
 
+def validate_config(config):
+    # PyTorch check.
+    if config["use_pytorch"]:
+        raise ValueError("DDPG does not support PyTorch yet! Use tf instead.")
+
+
 MARWILTrainer = build_trainer(
     name="MARWIL",
     default_config=DEFAULT_CONFIG,
     default_policy=MARWILPolicy,
-    make_policy_optimizer=make_optimizer)
+    validate_config=validate_config,
+    make_policy_optimizer=make_optimizer
+)

rllib/agents/pg/pg.py

@@ -9,8 +9,6 @@ DEFAULT_CONFIG = with_common_config({
     "num_workers": 0,
     # Learning rate.
     "lr": 0.0004,
-    # Use PyTorch as framework?
-    "use_pytorch": False
 })
 # __sphinx_doc_end__
 # yapf: enable

rllib/agents/ppo/ppo.py

@@ -142,14 +142,18 @@ def warn_about_bad_reward_scales(trainer, result):
 
 
 def validate_config(config):
+    # PyTorch check.
+    if config["use_pytorch"]:
+        raise ValueError("PPO does not support PyTorch yet! Use tf instead.")
     if config["entropy_coeff"] < 0:
         raise DeprecationWarning("entropy_coeff must be >= 0")
     if isinstance(config["entropy_coeff"], int):
         config["entropy_coeff"] = float(config["entropy_coeff"])
     if config["sgd_minibatch_size"] > config["train_batch_size"]:
         raise ValueError(
-            "Minibatch size {} must be <= train batch size {}.".format(
-                config["sgd_minibatch_size"], config["train_batch_size"]))
+            "Minibatch size {} must be <= train batch size {}.".
+            format(config["sgd_minibatch_size"], config["train_batch_size"])
+        )
     if config["batch_mode"] == "truncate_episodes" and not config["use_gae"]:
         raise ValueError(
             "Episode truncation is not supported without a value "

rllib/agents/trainer.py

@@ -146,6 +146,13 @@ COMMON_CONFIG = {
     # Log system resource metrics to results. This requires `psutil` to be
     # installed for sys stats, and `gputil` for GPU metrics.
     "log_sys_usage": True,
+
+    # === Framework Settings ===
+    # Use PyTorch (instead of tf). If using `rllib train`, this can also be
+    # enabled with the `--torch` flag.
+    # NOTE: Some agents may not support `torch` yet and throw an error.
+    "use_pytorch": False,
+
     # Enable TF eager execution (TF policies only). If using `rllib train`,
     # this can also be enabled with the `--eager` flag.
     "eager": False,

rllib/train.py

@@ -9,6 +9,11 @@ from ray.tune.config_parser import make_parser
 from ray.tune.result import DEFAULT_RESULTS_DIR
 from ray.tune.resources import resources_to_json
 from ray.tune.tune import _make_scheduler, run_experiments
+from ray.rllib.utils.framework import try_import_tf, try_import_torch
+
+# Try to import both backends for flag checking/warnings.
+tf = try_import_tf()
+torch, _ = try_import_torch()
 
 EXAMPLE_USAGE = """
 Training example via RLlib CLI:
@@ -92,6 +97,10 @@ def create_parser(parser_creator=None):
         "--resume",
         action="store_true",
         help="Whether to attempt to resume previous Tune experiments.")
+    parser.add_argument(
+        "--torch",
+        action="store_true",
+        help="Whether to use PyTorch (instead of tf) as the DL framework.")
     parser.add_argument(
         "--eager",
         action="store_true",
@@ -151,6 +160,8 @@ def run(args, parser):
             parser.error("the following arguments are required: --env")
         if args.eager:
             exp["config"]["eager"] = True
+        if args.torch:
+            exp["config"]["use_pytorch"] = True
         if args.v:
             exp["config"]["log_level"] = "INFO"
             verbose = 2