From c2cb5c2214f5427415f3fff770840cce6881aa28 Mon Sep 17 00:00:00 2001 From: Sven Mika Date: Tue, 7 Apr 2020 01:38:50 +0200 Subject: [PATCH] [RLlib] MARWIL torch. (#7836) * WIP. * WIP. * LINT. * Fix MARWIL so it can run with eager-mode. * LINT. --- rllib/BUILD | 29 ++++++- rllib/agents/marwil/__init__.py | 9 +- rllib/agents/marwil/marwil.py | 17 ++-- .../{marwil_policy.py => marwil_tf_policy.py} | 51 ++++++----- rllib/agents/marwil/marwil_torch_policy.py | 86 +++++++++++++++++++ rllib/agents/marwil/tests/test_marwil.py | 36 ++++++++ rllib/policy/torch_policy.py | 4 +- ...le-marwil.yaml => cartpole-marwil-tf.yaml} | 2 +- .../tuned_examples/cartpole-marwil-torch.yaml | 13 +++ 9 files changed, 209 insertions(+), 38 deletions(-) rename rllib/agents/marwil/{marwil_policy.py => marwil_tf_policy.py} (76%) create mode 100644 rllib/agents/marwil/marwil_torch_policy.py create mode 100644 rllib/agents/marwil/tests/test_marwil.py rename rllib/tuned_examples/{cartpole-marwil.yaml => cartpole-marwil-tf.yaml} (95%) create mode 100644 rllib/tuned_examples/cartpole-marwil-torch.yaml diff --git a/rllib/BUILD b/rllib/BUILD index fb813a328..977c78df4 100644 --- a/rllib/BUILD +++ b/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( diff --git a/rllib/agents/marwil/__init__.py b/rllib/agents/marwil/__init__.py index 606115be4..f901cf072 100644 --- a/rllib/agents/marwil/__init__.py +++ b/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", +] diff --git a/rllib/agents/marwil/marwil.py b/rllib/agents/marwil/marwil.py index ba96a14d8..3041041ab 100644 --- a/rllib/agents/marwil/marwil.py +++ b/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) diff --git a/rllib/agents/marwil/marwil_policy.py b/rllib/agents/marwil/marwil_tf_policy.py similarity index 76% rename from rllib/agents/marwil/marwil_policy.py rename to rllib/agents/marwil/marwil_tf_policy.py index 972d8f8ed..eab6f4ee9 100644 --- a/rllib/agents/marwil/marwil_policy.py +++ b/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( diff --git a/rllib/agents/marwil/marwil_torch_policy.py b/rllib/agents/marwil/marwil_torch_policy.py new file mode 100644 index 000000000..ea06b1b24 --- /dev/null +++ b/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]) diff --git a/rllib/agents/marwil/tests/test_marwil.py b/rllib/agents/marwil/tests/test_marwil.py new file mode 100644 index 000000000..c839841d2 --- /dev/null +++ b/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__])) diff --git a/rllib/policy/torch_policy.py b/rllib/policy/torch_policy.py index 6b896028f..cc77ab76f 100644 --- a/rllib/policy/torch_policy.py +++ b/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) diff --git a/rllib/tuned_examples/cartpole-marwil.yaml b/rllib/tuned_examples/cartpole-marwil-tf.yaml similarity index 95% rename from rllib/tuned_examples/cartpole-marwil.yaml rename to rllib/tuned_examples/cartpole-marwil-tf.yaml index 534d89745..e566f731a 100644 --- a/rllib/tuned_examples/cartpole-marwil.yaml +++ b/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: diff --git a/rllib/tuned_examples/cartpole-marwil-torch.yaml b/rllib/tuned_examples/cartpole-marwil-torch.yaml new file mode 100644 index 000000000..7b3d0fe99 --- /dev/null +++ b/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