[RLlib] MARWIL torch. (#7836)

* WIP.

* WIP.

* LINT.

* Fix MARWIL so it can run with eager-mode.

* LINT.

rllib/BUILD

@@ -115,6 +115,14 @@ py_test(
     srcs = ["agents/impala/tests/test_vtrace.py"]
 )
 
+# MARWILTrainer
+py_test(
+    name = "test_marwil",
+    tags = ["agents_dir"],
+    size = "small",
+    srcs = ["agents/marwil/tests/test_marwil.py"]
+)
+
 # PGTrainer
 py_test(
     name = "test_pg",
@@ -573,8 +581,9 @@ py_test(
 # MARWIL
 
 py_test(
-    name = "test_marwil_cartpole_v0",
-    main = "train.py", srcs = ["train.py"],
+    name = "test_marwil_cartpole_v0_tf",
+    main = "train.py",
+    srcs = ["train.py"],
     tags = ["quick_train", "external_files"],
     size = "small",
     # Include the json data file.
@@ -587,6 +596,22 @@ py_test(
         ]
 )
 
+py_test(
+    name = "test_marwil_cartpole_v0_torch",
+    main = "train.py",
+    srcs = ["train.py"],
+    tags = ["quick_train", "external_files"],
+    size = "small",
+    # Include the json data file.
+    data = glob(["tests/data/cartpole_small/**"]),
+    args = [
+        "--env", "CartPole-v0",
+        "--run", "MARWIL",
+        "--stop", "'{\"training_iteration\": 1}'",
+        "--config", "'{\"use_pytorch\": true, \"input\": \"tests/data/cartpole_small\", \"learning_starts\": 0, \"input_evaluation\": [\"wis\", \"is\"], \"shuffle_buffer_size\": 10}'"
+        ]
+)
+
 # PG
 
 py_test(

rllib/agents/marwil/__init__.py

@@ -1,3 +1,10 @@
 from ray.rllib.agents.marwil.marwil import MARWILTrainer, DEFAULT_CONFIG
+from ray.rllib.agents.marwil.marwil_tf_policy import MARWILTFPolicy
+from ray.rllib.agents.marwil.marwil_torch_policy import MARWILTorchPolicy
 
-__all__ = ["MARWILTrainer", "DEFAULT_CONFIG"]
+__all__ = [
+    "DEFAULT_CONFIG",
+    "MARWILTFPolicy",
+    "MARWILTorchPolicy",
+    "MARWILTrainer",
+]

rllib/agents/marwil/marwil.py

@@ -1,6 +1,6 @@
 from ray.rllib.agents.trainer import with_common_config
 from ray.rllib.agents.trainer_template import build_trainer
-from ray.rllib.agents.marwil.marwil_policy import MARWILTFPolicy
+from ray.rllib.agents.marwil.marwil_tf_policy import MARWILTFPolicy
 from ray.rllib.optimizers import SyncBatchReplayOptimizer
 
 # yapf: disable
@@ -30,6 +30,8 @@ DEFAULT_CONFIG = with_common_config({
     "learning_starts": 0,
     # === Parallelism ===
     "num_workers": 0,
+    # Use PyTorch as framework?
+    "use_pytorch": False
 })
 # __sphinx_doc_end__
 # yapf: enable
@@ -44,15 +46,18 @@ 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.")
+def get_policy_class(config):
+    if config.get("use_pytorch") is True:
+        from ray.rllib.agents.marwil.marwil_torch_policy import \
+            MARWILTorchPolicy
+        return MARWILTorchPolicy
+    else:
+        return MARWILTFPolicy
 
 
 MARWILTrainer = build_trainer(
     name="MARWIL",
     default_config=DEFAULT_CONFIG,
     default_policy=MARWILTFPolicy,
-    validate_config=validate_config,
+    get_policy_class=get_policy_class,
     make_policy_optimizer=make_optimizer)

rllib/agents/marwil/marwil_tf_policy.py

@@ -1,7 +1,3 @@
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import ray
 from ray.rllib.policy.sample_batch import SampleBatch
 from ray.rllib.utils.explained_variance import explained_variance
@@ -14,7 +10,7 @@ from ray.rllib.utils import try_import_tf
 tf = try_import_tf()
 
 
-class ValueNetworkMixin(object):
+class ValueNetworkMixin:
     def __init__(self):
         @make_tf_callable(self.get_session())
         def value(ob, prev_action, prev_reward, *state):
@@ -30,31 +26,27 @@ class ValueNetworkMixin(object):
         self._value = value
 
 
-class ValueLoss(object):
+class ValueLoss:
     def __init__(self, state_values, cumulative_rewards):
         self.loss = 0.5 * tf.reduce_mean(
             tf.square(state_values - cumulative_rewards))
 
 
-class ReweightedImitationLoss(object):
-    def __init__(self, state_values, cumulative_rewards, actions, action_dist,
-                 beta):
-        ma_adv_norm = tf.get_variable(
-            name="moving_average_of_advantage_norm",
-            dtype=tf.float32,
-            initializer=100.0,
-            trainable=False)
+class ReweightedImitationLoss:
+    def __init__(self, policy, state_values, cumulative_rewards, actions,
+                 action_dist, beta):
         # advantage estimation
         adv = cumulative_rewards - state_values
         # update averaged advantage norm
         update_adv_norm = tf.assign_add(
-            ref=ma_adv_norm,
-            value=1e-6 * (tf.reduce_mean(tf.square(adv)) - ma_adv_norm))
+            ref=policy._ma_adv_norm,
+            value=1e-6 *
+            (tf.reduce_mean(tf.square(adv)) - policy._ma_adv_norm))
 
         # exponentially weighted advantages
         with tf.control_dependencies([update_adv_norm]):
             exp_advs = tf.exp(
-                beta * tf.divide(adv, 1e-8 + tf.sqrt(ma_adv_norm)))
+                beta * tf.divide(adv, 1e-8 + tf.sqrt(policy._ma_adv_norm)))
 
         # log\pi_\theta(a|s)
         logprobs = action_dist.logp(actions)
@@ -87,24 +79,24 @@ def postprocess_advantages(policy,
         use_critic=False)
 
 
-class MARWILLoss(object):
-    def __init__(self, state_values, action_dist, actions, advantages,
+class MARWILLoss:
+    def __init__(self, policy, state_values, action_dist, actions, advantages,
                  vf_loss_coeff, beta):
 
         self.v_loss = self._build_value_loss(state_values, advantages)
-        self.p_loss = self._build_policy_loss(state_values, advantages,
+        self.p_loss = self._build_policy_loss(policy, state_values, advantages,
                                               actions, action_dist, beta)
 
         self.total_loss = self.p_loss.loss + vf_loss_coeff * self.v_loss.loss
-        self.explained_variance = tf.reduce_mean(
-            explained_variance(advantages, state_values))
+        explained_var = explained_variance(advantages, state_values)
+        self.explained_variance = tf.reduce_mean(explained_var)
 
     def _build_value_loss(self, state_values, cum_rwds):
         return ValueLoss(state_values, cum_rwds)
 
-    def _build_policy_loss(self, state_values, cum_rwds, actions, action_dist,
-                           beta):
-        return ReweightedImitationLoss(state_values, cum_rwds, actions,
+    def _build_policy_loss(self, policy, state_values, cum_rwds, actions,
+                           action_dist, beta):
+        return ReweightedImitationLoss(policy, state_values, cum_rwds, actions,
                                        action_dist, beta)
 
 
@@ -113,7 +105,7 @@ def marwil_loss(policy, model, dist_class, train_batch):
     action_dist = dist_class(model_out, model)
     state_values = model.value_function()
 
-    policy.loss = MARWILLoss(state_values, action_dist,
+    policy.loss = MARWILLoss(policy, state_values, action_dist,
                              train_batch[SampleBatch.ACTIONS],
                              train_batch[Postprocessing.ADVANTAGES],
                              policy.config["vf_coeff"], policy.config["beta"])
@@ -132,6 +124,13 @@ def stats(policy, train_batch):
 
 def setup_mixins(policy, obs_space, action_space, config):
     ValueNetworkMixin.__init__(policy)
+    # Set up a tf-var for the moving avg (do this here to make it work with
+    # eager mode).
+    policy._ma_adv_norm = tf.get_variable(
+        name="moving_average_of_advantage_norm",
+        dtype=tf.float32,
+        initializer=100.0,
+        trainable=False)
 
 
 MARWILTFPolicy = build_tf_policy(

rllib/agents/marwil/marwil_torch_policy.py

@@ -0,0 +1,86 @@
+import ray
+from ray.rllib.agents.marwil.marwil_tf_policy import postprocess_advantages
+from ray.rllib.evaluation.postprocessing import Postprocessing
+from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.policy.torch_policy_template import build_torch_policy
+from ray.rllib.utils.explained_variance import explained_variance
+from ray.rllib.utils.framework import try_import_torch
+
+torch, _ = try_import_torch()
+
+
+class ValueNetworkMixin:
+    def __init__(self):
+        def value(ob, prev_action, prev_reward, *state):
+            model_out, _ = self.model({
+                SampleBatch.CUR_OBS: torch.Tensor([ob]).to(self.device),
+                SampleBatch.PREV_ACTIONS: torch.Tensor([prev_action]).to(
+                    self.device),
+                SampleBatch.PREV_REWARDS: torch.Tensor([prev_reward]).to(
+                    self.device),
+                "is_training": False,
+            }, [torch.Tensor([s]).to(self.device) for s in state],
+                                      torch.Tensor([1]).to(self.device))
+            return self.model.value_function()[0]
+
+        self._value = value
+
+
+def marwil_loss(policy, model, dist_class, train_batch):
+    model_out, _ = model.from_batch(train_batch)
+    action_dist = dist_class(model_out, model)
+    state_values = model.value_function()
+    advantages = train_batch[Postprocessing.ADVANTAGES]
+    actions = train_batch[SampleBatch.ACTIONS]
+
+    # Value loss.
+    policy.v_loss = 0.5 * torch.mean(torch.pow(state_values - advantages, 2.0))
+
+    # Policy loss.
+    # Advantage estimation.
+    adv = advantages - state_values
+    # Update averaged advantage norm.
+    policy.ma_adv_norm.add_(
+        1e-6 * (torch.mean(torch.pow(adv, 2.0)) - policy.ma_adv_norm))
+    # #xponentially weighted advantages.
+    exp_advs = torch.exp(policy.config["beta"] *
+                         (adv / (1e-8 + torch.pow(policy.ma_adv_norm, 0.5))))
+    # log\pi_\theta(a|s)
+    logprobs = action_dist.logp(actions)
+    policy.p_loss = -1.0 * torch.mean(exp_advs.detach() * logprobs)
+
+    # Combine both losses.
+    policy.total_loss = policy.p_loss + policy.config["vf_coeff"] * \
+        policy.v_loss
+    explained_var = explained_variance(
+        advantages, state_values, framework="torch")
+    policy.explained_variance = torch.mean(explained_var)
+
+    return policy.total_loss
+
+
+def stats(policy, train_batch):
+    return {
+        "policy_loss": policy.p_loss,
+        "vf_loss": policy.v_loss,
+        "total_loss": policy.total_loss,
+        "vf_explained_var": policy.explained_variance,
+    }
+
+
+def setup_mixins(policy, obs_space, action_space, config):
+    # Create a var.
+    policy.ma_adv_norm = torch.tensor(
+        [100.0], dtype=torch.float32, requires_grad=False)
+    # Setup Value branch of our NN.
+    ValueNetworkMixin.__init__(policy)
+
+
+MARWILTorchPolicy = build_torch_policy(
+    name="MARWILTorchPolicy",
+    loss_fn=marwil_loss,
+    get_default_config=lambda: ray.rllib.agents.marwil.marwil.DEFAULT_CONFIG,
+    stats_fn=stats,
+    postprocess_fn=postprocess_advantages,
+    after_init=setup_mixins,
+    mixins=[ValueNetworkMixin])

rllib/agents/marwil/tests/test_marwil.py

@@ -0,0 +1,36 @@
+import unittest
+
+import ray
+import ray.rllib.agents.marwil as marwil
+from ray.rllib.utils.framework import try_import_tf
+from ray.rllib.utils.test_utils import framework_iterator
+
+tf = try_import_tf()
+
+
+class TestMARWIL(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        ray.init()
+
+    @classmethod
+    def tearDownClass(cls):
+        ray.shutdown()
+
+    def test_marwil_compilation(self):
+        """Test whether a MARWILTrainer can be built with all frameworks."""
+        config = marwil.DEFAULT_CONFIG.copy()
+        config["num_workers"] = 0  # Run locally.
+        num_iterations = 2
+
+        # Test for all frameworks.
+        for _ in framework_iterator(config):
+            trainer = marwil.MARWILTrainer(config=config, env="CartPole-v0")
+            for i in range(num_iterations):
+                trainer.train()
+
+
+if __name__ == "__main__":
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))

rllib/policy/torch_policy.py

@@ -191,9 +191,9 @@ class TorchPolicy(Policy):
                 SampleBatch.CUR_OBS: obs_batch,
                 SampleBatch.ACTIONS: actions
             })
-            if prev_action_batch:
+            if prev_action_batch is not None:
                 input_dict[SampleBatch.PREV_ACTIONS] = prev_action_batch
-            if prev_reward_batch:
+            if prev_reward_batch is not None:
                 input_dict[SampleBatch.PREV_REWARDS] = prev_reward_batch
             seq_lens = torch.ones(len(obs_batch), dtype=torch.int32)
 

rllib/tuned_examples/cartpole-marwil-tf.yaml

@@ -2,7 +2,7 @@
 # $ ./train.py --run=PPO --env=CartPole-v0 \
 #      --stop='{"timesteps_total": 50000}' \
 #      --config='{"output": "/tmp/out", "batch_mode": "complete_episodes"}'
-cartpole-marwil:
+cartpole-marwil-tf:
     env: CartPole-v0
     run: MARWIL
     stop:

rllib/tuned_examples/cartpole-marwil-torch.yaml

@@ -0,0 +1,13 @@
+# To generate training data, first run:
+# $ ./train.py --run=PPO --env=CartPole-v0 \
+#      --stop='{"timesteps_total": 50000}' \
+#      --config='{"use_pytorch": true, "output": "/tmp/out", "batch_mode": "complete_episodes"}'
+cartpole-marwil-torch:
+    env: CartPole-v0
+    run: MARWIL
+    stop:
+        timesteps_total: 500000
+    config:
+        beta:
+            grid_search: [0, 1]   # compare IL (beta=0) vs MARWIL
+        input: /tmp/out