From 14405b90d5457863d71168c613b4961d34f19cc5 Mon Sep 17 00:00:00 2001 From: Sven Mika Date: Tue, 16 Jun 2020 09:01:20 +0200 Subject: [PATCH] [RLlib] Prototype of a DynaTrainer (for env dynamics learning in upcoming MBMPO algo). (#8860) --- rllib/BUILD | 8 ++ rllib/agents/dyna/__init__.py | 10 +++ rllib/agents/dyna/dyna.py | 101 +++++++++++++++++++++++ rllib/agents/dyna/dyna_torch_model.py | 58 +++++++++++++ rllib/agents/dyna/dyna_torch_policy.py | 94 +++++++++++++++++++++ rllib/agents/dyna/tests/test_dyna.py | 72 ++++++++++++++++ rllib/examples/env/stateless_cartpole.py | 2 +- rllib/models/catalog.py | 55 ++++++------ rllib/policy/torch_policy.py | 9 ++ rllib/policy/torch_policy_template.py | 4 +- 10 files changed, 387 insertions(+), 26 deletions(-) create mode 100644 rllib/agents/dyna/__init__.py create mode 100644 rllib/agents/dyna/dyna.py create mode 100644 rllib/agents/dyna/dyna_torch_model.py create mode 100644 rllib/agents/dyna/dyna_torch_policy.py create mode 100644 rllib/agents/dyna/tests/test_dyna.py diff --git a/rllib/BUILD b/rllib/BUILD index 795452cf7..0674540e9 100644 --- a/rllib/BUILD +++ b/rllib/BUILD @@ -433,6 +433,14 @@ py_test( srcs = ["agents/dqn/tests/test_simple_q.py"] ) +# DYNATrainer +py_test( + name = "test_dyna", + tags = ["agents_dir"], + size = "small", + srcs = ["agents/dyna/tests/test_dyna.py"] +) + # ES py_test( name = "test_es", diff --git a/rllib/agents/dyna/__init__.py b/rllib/agents/dyna/__init__.py new file mode 100644 index 000000000..5bb73fb3b --- /dev/null +++ b/rllib/agents/dyna/__init__.py @@ -0,0 +1,10 @@ +from ray.rllib.agents.dyna.dyna import DYNATrainer, DEFAULT_CONFIG +from ray.rllib.agents.dyna.dyna_torch_policy import dyna_torch_loss, \ + DYNATorchPolicy + +__all__ = [ + "dyna_torch_loss", + "DEFAULT_CONFIG", + "DYNATorchPolicy", + "DYNATrainer", +] diff --git a/rllib/agents/dyna/dyna.py b/rllib/agents/dyna/dyna.py new file mode 100644 index 000000000..405128c15 --- /dev/null +++ b/rllib/agents/dyna/dyna.py @@ -0,0 +1,101 @@ +import logging + +from ray.rllib.agents.trainer import with_common_config +from ray.rllib.agents.trainer_template import build_trainer + +logger = logging.getLogger(__name__) + +# yapf: disable +# __sphinx_doc_begin__ +DEFAULT_CONFIG = with_common_config({ + # Default Trainer setting overrides. + "num_workers": 1, + "num_envs_per_worker": 1, + + # The size of an entire epoch (for supervised learning the dynamics). + # The train-batch will be split into training and validation sets according + # to `training_set_ratio`, then n epochs (with minibatch + # size=`sgd_minibatch_size`) will be trained until the sliding average + # of the validation performance decreases. + "train_batch_size": 10000, + "sgd_minibatch_size": 500, + "rollout_fragment_length": 200, + # Learning rate for the dynamics optimizer. + "lr": 0.0003, + + # Fraction of the entire data that should be used for training the dynamics + # model. The validation fraction is 1.0 - `training_set_ratio`. Training of + # a dynamics model over n some epochs (1 epoch = entire training set) stops + # when the validation set's performance starts to decrease. + "train_set_ratio": 0.8, + + # The exploration strategy to apply on top of the (acting) policy. + # TODO: (sven) Use random for testing purposes for now. + "exploration_config": {"type": "Random"}, + + # Whether to predict the action that lead from obs(t) to obs(t+1), instead + # of predicting obs(t+1). + "predict_action": False, + + # Whether the dynamics model should predict the reward, given obs(t)+a(t). + # NOTE: Only supported if `predict_action`=False. + "predict_reward": False, + + # Whether to use the same network for predicting rewards than for + # predicting the next observation. + "reward_share_layers": True, + + # TODO: (sven) figure out API to query the latent space vector given + # some observation (not needed for MBMPO). + "learn_latent_space": False, + + # Whether to predict `obs(t+1) - obs(t)` instead of `obs(t+1)` directly. + # NOTE: This only works for 1D Box observation spaces, e.g. Box(5,) and + # if `predict_action`=False. + "predict_obs_delta": True, + # TODO: loss function types: neg_log_llh, etc..? + "loss_function": "l2", + + # Config for the dynamics learning model architecture. + "dynamics_model": { + "fcnet_hiddens": [512, 512], + "fcnet_activation": "relu", + }, + + # TODO: (sven) allow for having a default model config over many + # sub-models: e.g. "model": {"ModelA": {[default_config]}, + # "ModelB": [default_config]} +}) +# __sphinx_doc_end__ +# yapf: enable + + +def validate_config(config): + if config["train_set_ratio"] <= 0.0 or \ + config["train_set_ratio"] >= 1.0: + raise ValueError("`train_set_ratio` must be within (0.0, 1.0)!") + if config["predict_action"] or config["predict_reward"]: + raise ValueError( + "`predict_action`=True or `predict_reward`=True not supported " + "yet!") + if config["learn_latent_space"]: + raise ValueError("`learn_latent_space` not supported yet!") + if config["loss_function"] != "l2": + raise ValueError("`loss_function` other than 'l2' not supported yet!") + + +def get_policy_class(config): + if config["framework"] == "torch": + from ray.rllib.agents.dyna.dyna_torch_policy import DYNATorchPolicy + return DYNATorchPolicy + else: + raise ValueError("tf not supported yet!") + + +DYNATrainer = build_trainer( + name="DYNA", + default_policy=None, + get_policy_class=get_policy_class, + default_config=DEFAULT_CONFIG, + validate_config=validate_config, +) diff --git a/rllib/agents/dyna/dyna_torch_model.py b/rllib/agents/dyna/dyna_torch_model.py new file mode 100644 index 000000000..6c748b7aa --- /dev/null +++ b/rllib/agents/dyna/dyna_torch_model.py @@ -0,0 +1,58 @@ +import gym +from gym.spaces import Discrete + +from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 +from ray.rllib.utils.framework import try_import_torch + +torch, nn = try_import_torch() + + +class DYNATorchModel(TorchModelV2, nn.Module): + """Extension of standard TorchModelV2 for Env dynamics learning. + + Data flow: + obs.cat(action) -> forward() -> next_obs|next_obs_delta + get_next_state(obs, action) -> next_obs|next_obs_delta + + Note that this class by itself is not a valid model unless you + implement forward() in a subclass. + """ + + def __init__(self, obs_space, action_space, num_outputs, model_config, + name): + """Initializes a DYNATorchModel object. + """ + + nn.Module.__init__(self) + # Construct the wrapped model handing it a concat'd observation and + # action space as "input_space" and our obs_space as "output_space". + # TODO: (sven) get rid of these restrictions on obs/action spaces. + assert isinstance(action_space, Discrete) + input_space = gym.spaces.Box( + obs_space.low[0], + obs_space.high[0], + shape=(obs_space.shape[0] + action_space.n, )) + super(DYNATorchModel, self).__init__(input_space, action_space, + num_outputs, model_config, name) + + def get_next_observation(self, observations, actions): + """Returns a next obs prediction given current observation and action. + + This implements p^(s'|s, a). With p being the environment dynamics. + + Arguments: + observations (Tensor): The current observation Tensor. + actions (Tensor): The actions taken in `observations`. + + Returns: + TensorType: The predicted next observations. + """ + + # One-hot the actions. + actions_flat = nn.functional.one_hot( + actions, num_classes=self.action_space.n).float() + # Push through our underlying Model. + next_obs, _ = self.forward({ + "obs_flat": torch.cat([observations, actions_flat], -1) + }, [], None) + return next_obs diff --git a/rllib/agents/dyna/dyna_torch_policy.py b/rllib/agents/dyna/dyna_torch_policy.py new file mode 100644 index 000000000..c2ca52aeb --- /dev/null +++ b/rllib/agents/dyna/dyna_torch_policy.py @@ -0,0 +1,94 @@ +import gym +import logging + +import ray +from ray.rllib.agents.dyna.dyna_torch_model import DYNATorchModel +from ray.rllib.models.catalog import ModelCatalog +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.torch_policy_template import build_torch_policy +from ray.rllib.utils import try_import_torch + +torch, nn = try_import_torch() + +logger = logging.getLogger(__name__) + + +def make_model_and_dist(policy, obs_space, action_space, config): + # Get the output distribution class for predicting rewards and next-obs. + policy.distr_cls_next_obs, num_outputs = ModelCatalog.get_action_dist( + obs_space, config, dist_type="deterministic", framework="torch") + if config["predict_reward"]: + # TODO: (sven) implement reward prediction. + _ = ModelCatalog.get_action_dist( + gym.spaces.Box(float("-inf"), float("inf"), ()), + config, + dist_type="") + + # Build one dynamics model if we are a Worker. + # If we are the main MAML learner, build n (num_workers) dynamics Models + # for being able to create checkpoints for the current state of training. + policy.dynamics_model = ModelCatalog.get_model_v2( + obs_space, + action_space, + num_outputs=num_outputs, + model_config=config["dynamics_model"], + framework="torch", + name="dynamics_model", + model_interface=DYNATorchModel, + ) + + action_dist, num_outputs = ModelCatalog.get_action_dist( + action_space, config, dist_type="deterministic", framework="torch") + # Create the pi-model and register it with the Policy. + policy.pi = ModelCatalog.get_model_v2( + obs_space, + action_space, + num_outputs=num_outputs, + model_config=config["model"], + framework="torch", + name="policy_model", + ) + + return policy.pi, action_dist + + +def dyna_torch_loss(policy, model, dist_class, train_batch): + # Split batch into train and validation sets according to + # `train_set_ratio`. + predicted_next_state_deltas = \ + policy.dynamics_model.get_next_observation( + train_batch[SampleBatch.CUR_OBS], train_batch[SampleBatch.ACTIONS]) + labels = train_batch[SampleBatch.NEXT_OBS] - train_batch[SampleBatch. + CUR_OBS] + loss = torch.pow( + torch.sum( + torch.pow(labels - predicted_next_state_deltas, 2.0), dim=-1), 0.5) + batch_size = int(loss.shape[0]) + train_set_size = int(batch_size * policy.config["train_set_ratio"]) + train_loss, validation_loss = \ + torch.split(loss, (train_set_size, batch_size - train_set_size), dim=0) + policy.dynamics_train_loss = torch.mean(train_loss) + policy.dynamics_validation_loss = torch.mean(validation_loss) + return policy.dynamics_train_loss + + +def stats_fn(policy, train_batch): + return { + "dynamics_train_loss": policy.dynamics_train_loss, + "dynamics_validation_loss": policy.dynamics_validation_loss, + } + + +def torch_optimizer(policy, config): + return torch.optim.Adam( + policy.dynamics_model.parameters(), lr=config["lr"]) + + +DYNATorchPolicy = build_torch_policy( + name="DYNATorchPolicy", + loss_fn=dyna_torch_loss, + get_default_config=lambda: ray.rllib.agents.dyna.dyna.DEFAULT_CONFIG, + stats_fn=stats_fn, + optimizer_fn=torch_optimizer, + make_model_and_action_dist=make_model_and_dist, +) diff --git a/rllib/agents/dyna/tests/test_dyna.py b/rllib/agents/dyna/tests/test_dyna.py new file mode 100644 index 000000000..d9e36fe5a --- /dev/null +++ b/rllib/agents/dyna/tests/test_dyna.py @@ -0,0 +1,72 @@ +import copy +import gym +import numpy as np +import unittest + +import ray +import ray.rllib.agents.dyna as dyna +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.test_utils import check_compute_single_action, \ + framework_iterator + +torch, _ = try_import_torch() + + +class TestDYNA(unittest.TestCase): + @classmethod + def setUpClass(cls): + ray.init(local_mode=True) + + @classmethod + def tearDownClass(cls): + ray.shutdown() + + def test_dyna_compilation(self): + """Test whether a DYNATrainer can be built with both frameworks.""" + config = copy.deepcopy(dyna.DEFAULT_CONFIG) + config["num_workers"] = 1 + config["train_batch_size"] = 1000 + num_iterations = 30 + env = "CartPole-v0" + test_env = gym.make(env) + + for _ in framework_iterator(config, frameworks="torch"): + trainer = dyna.DYNATrainer(config=config, env=env) + policy = trainer.get_policy() + # Do n supervised epochs, each over `train_batch_size`. + # Ignore validation loss here as a stopping criteria. + for i in range(num_iterations): + info = trainer.train()["info"]["learner"]["default_policy"] + print("SL iteration: {}".format(i)) + print("train loss {}".format(info["dynamics_train_loss"])) + print("validation loss {}".format( + info["dynamics_validation_loss"])) + # Check, whether normal action stepping works with DYNA's policy. + # Note that DYNA does not train its Policy. It must be pushed + # down from the main model-based algo from time to time. + check_compute_single_action(trainer) + + # Check, whether env dynamics were actually learnt - more or less. + obs = test_env.reset() + for _ in range(10): + action = trainer.compute_action(obs) + obs = torch.from_numpy(np.array([obs])).float() + # Make the prediction over the next state (deterministic delta + # like in MBMPO). + predicted_next_obs_delta = \ + policy.dynamics_model.get_next_observation( + obs, + torch.from_numpy(np.array([action]))) + predicted_next_obs = obs + predicted_next_obs_delta + obs, _, done, _ = test_env.step(action) + self.assertLess( + np.sum(obs - predicted_next_obs.detach().numpy()), 0.05) + # Reset if done. + if done: + obs = test_env.reset() + + +if __name__ == "__main__": + import pytest + import sys + sys.exit(pytest.main(["-v", __file__])) diff --git a/rllib/examples/env/stateless_cartpole.py b/rllib/examples/env/stateless_cartpole.py index dd720414e..67c4902e1 100644 --- a/rllib/examples/env/stateless_cartpole.py +++ b/rllib/examples/env/stateless_cartpole.py @@ -1,7 +1,7 @@ -import math import gym from gym import spaces from gym.utils import seeding +import math import numpy as np diff --git a/rllib/models/catalog.py b/rllib/models/catalog.py index 5e2135f4c..e6f7cc526 100644 --- a/rllib/models/catalog.py +++ b/rllib/models/catalog.py @@ -125,14 +125,18 @@ class ModelCatalog: Args: action_space (Space): Action space of the target gym env. config (Optional[dict]): Optional model config. - dist_type (Optional[str]): Identifier of the action distribution. + dist_type (Optional[str]): Identifier of the action distribution + interpreted as a hint. framework (str): One of "tf", "tfe", or "torch". kwargs (dict): Optional kwargs to pass on to the Distribution's constructor. Returns: - dist_class (ActionDistribution): Python class of the distribution. - dist_dim (int): The size of the input vector to the distribution. + Tuple: + - dist_class (ActionDistribution): Python class of the + distribution. + - dist_dim (int): The size of the input vector to the + distribution. """ dist = None @@ -523,6 +527,30 @@ class ModelCatalog: return wrapper + @staticmethod + def _get_v2_model_class(input_space, model_config, framework="tf"): + if framework == "torch": + from ray.rllib.models.torch.fcnet import (FullyConnectedNetwork as + FCNet) + from ray.rllib.models.torch.visionnet import (VisionNetwork as + VisionNet) + else: + from ray.rllib.models.tf.fcnet import \ + FullyConnectedNetwork as FCNet + from ray.rllib.models.tf.visionnet import \ + VisionNetwork as VisionNet + + # Discrete/1D obs-spaces. + if isinstance(input_space, gym.spaces.Discrete) or \ + len(input_space.shape) <= 2: + return FCNet + # Default Conv2D net. + else: + return VisionNet + + # ------------------- + # DEPRECATED METHODS. + # ------------------- @staticmethod def get_model(input_dict, obs_space, @@ -581,27 +609,6 @@ class ModelCatalog: return FullyConnectedNetwork(input_dict, obs_space, action_space, num_outputs, options) - @staticmethod - def _get_v2_model_class(obs_space, model_config, framework="tf"): - if framework == "torch": - from ray.rllib.models.torch.fcnet import (FullyConnectedNetwork as - FCNet) - from ray.rllib.models.torch.visionnet import (VisionNetwork as - VisionNet) - else: - from ray.rllib.models.tf.fcnet import \ - FullyConnectedNetwork as FCNet - from ray.rllib.models.tf.visionnet import \ - VisionNetwork as VisionNet - - # Discrete/1D obs-spaces. - if isinstance(obs_space, gym.spaces.Discrete) or \ - len(obs_space.shape) <= 2: - return FCNet - # Default Conv2D net. - else: - return VisionNet - @staticmethod def get_torch_model(obs_space, num_outputs, diff --git a/rllib/policy/torch_policy.py b/rllib/policy/torch_policy.py index c13ff5f8a..eac18200d 100644 --- a/rllib/policy/torch_policy.py +++ b/rllib/policy/torch_policy.py @@ -1,6 +1,8 @@ +import functools import numpy as np import time +from ray.rllib.models.torch.torch_action_dist import TorchDistributionWrapper from ray.rllib.policy.policy import Policy, LEARNER_STATS_KEY from ray.rllib.policy.sample_batch import SampleBatch from ray.rllib.policy.rnn_sequencing import pad_batch_to_sequences_of_same_size @@ -149,6 +151,13 @@ class TorchPolicy(Policy): dist_class = self.dist_class dist_inputs, state_out = self.model( input_dict, state_batches, seq_lens) + if not (isinstance(dist_class, functools.partial) + or issubclass(dist_class, TorchDistributionWrapper)): + raise ValueError( + "`dist_class` ({}) not a TorchDistributionWrapper " + "subclass! Make sure your `action_distribution_fn` or " + "`make_model_and_action_dist` return a correct " + "distribution class.".format(dist_class.__name__)) action_dist = dist_class(dist_inputs, self.model) # Get the exploration action from the forward results. diff --git a/rllib/policy/torch_policy_template.py b/rllib/policy/torch_policy_template.py index e629f5c3c..7be5855e2 100644 --- a/rllib/policy/torch_policy_template.py +++ b/rllib/policy/torch_policy_template.py @@ -99,7 +99,9 @@ def build_torch_policy(name, # Model is customized (use default action dist class). if make_model: - assert make_model_and_action_dist is None + assert make_model_and_action_dist is None, \ + "Either `make_model` or `make_model_and_action_dist`" \ + " must be None!" self.model = make_model(self, obs_space, action_space, config) dist_class, _ = ModelCatalog.get_action_dist( action_space, self.config["model"], framework="torch")