[RLlib] Implementation of "Model-based Meta Policy Optimization" (MB MPO) (#9409)

rllib/agents/maml/maml_torch_policy.py

@@ -71,7 +71,8 @@ def PPOLoss(dist_class,
     entropy_loss = torch.mean(entropy_loss(pi_new_dist))
 
     total_loss = -surr_loss + cur_kl_coeff * kl_loss
-    total_loss += vf_loss_coeff * vf_loss - entropy_coeff * entropy_loss
+    total_loss += vf_loss_coeff * vf_loss
+    total_loss -= entropy_coeff * entropy_loss
     return total_loss, surr_loss, kl_loss, vf_loss, entropy_loss
 
 
@@ -108,7 +109,6 @@ class WorkerLoss(object):
             vf_clip_param=vf_clip_param,
             vf_loss_coeff=vf_loss_coeff,
             clip_loss=clip_loss)
-        print("Worker Loss: ", self.loss)
 
 
 # This is the Meta-Update computation graph for main (meta-update step)
@@ -303,8 +303,7 @@ class MAMLLoss(object):
             loss,
             inputs=model.parameters(),
             create_graph=True,
-            retain_graph=True,
-            only_inputs=True)
+            allow_unused=True)
         adapted_vars = {}
         for i, tup in enumerate(network_vars.items()):
             name, var = tup

rllib/agents/mbmpo/__init__.py

@@ -0,0 +1,6 @@
+from ray.rllib.agents.mbmpo.mbmpo import MBMPOTrainer, DEFAULT_CONFIG
+
+__all__ = [
+    "MBMPOTrainer",
+    "DEFAULT_CONFIG",
+]

rllib/agents/mbmpo/mbmpo.py

@@ -0,0 +1,389 @@
+import logging
+
+import numpy as np
+import ray
+from ray.rllib.utils.sgd import standardized
+from ray.rllib.agents import with_common_config
+from ray.rllib.agents.mbmpo.mbmpo_torch_policy import MBMPOTorchPolicy
+from ray.rllib.agents.trainer_template import build_trainer
+from ray.rllib.evaluation.metrics import get_learner_stats
+from ray.rllib.execution.common import STEPS_SAMPLED_COUNTER, \
+    STEPS_TRAINED_COUNTER, LEARNER_INFO, _get_shared_metrics
+from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.execution.metric_ops import CollectMetrics
+from ray.util.iter import from_actors
+from ray.rllib.agents.mbmpo.model_ensemble import DynamicsEnsembleCustomModel
+from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID
+from ray.rllib.utils.torch_ops import convert_to_torch_tensor
+from ray.rllib.evaluation.metrics import collect_episodes
+from ray.rllib.agents.mbmpo.model_vector_env import custom_model_vector_env
+from ray.rllib.evaluation.metrics import collect_metrics
+from ray.rllib.agents.mbmpo.utils import calculate_gae_advantages
+
+logger = logging.getLogger(__name__)
+
+# yapf: disable
+# __sphinx_doc_begin__
+DEFAULT_CONFIG = with_common_config({
+    # If true, use the Generalized Advantage Estimator (GAE)
+    # with a value function, see https://arxiv.org/pdf/1506.02438.pdf.
+    "use_gae": True,
+    # GAE(lambda) parameter
+    "lambda": 1.0,
+    # Initial coefficient for KL divergence
+    "kl_coeff": 0.0005,
+    # Size of batches collected from each worker
+    "rollout_fragment_length": 200,
+    # Stepsize of SGD
+    "lr": 1e-3,
+    # Share layers for value function
+    "vf_share_layers": False,
+    # Coefficient of the value function loss
+    "vf_loss_coeff": 0.5,
+    # Coefficient of the entropy regularizer
+    "entropy_coeff": 0.0,
+    # PPO clip parameter
+    "clip_param": 0.5,
+    # Clip param for the value function. Note that this is sensitive to the
+    # scale of the rewards. If your expected V is large, increase this.
+    "vf_clip_param": 10.0,
+    # If specified, clip the global norm of gradients by this amount
+    "grad_clip": None,
+    # Target value for KL divergence
+    "kl_target": 0.01,
+    # Whether to rollout "complete_episodes" or "truncate_episodes"
+    "batch_mode": "complete_episodes",
+    # Which observation filter to apply to the observation
+    "observation_filter": "NoFilter",
+    # Number of Inner adaptation steps for the MAML algorithm
+    "inner_adaptation_steps": 1,
+    # Number of MAML steps per meta-update iteration (PPO steps)
+    "maml_optimizer_steps": 8,
+    # Inner Adaptation Step size
+    "inner_lr": 1e-3,
+    # Horizon of Environment (200 in MB-MPO paper)
+    "horizon": 200,
+    # Dynamics Ensemble Hyperparameters
+    "dynamics_model": {
+        "custom_model": DynamicsEnsembleCustomModel,
+        # Number of Transition-Dynamics Models for Ensemble
+        "ensemble_size": 5,
+        # Hidden Layers for Model Ensemble
+        "fcnet_hiddens": [512, 512],
+        # Model Learning Rate
+        "lr": 1e-3,
+        # Max number of training epochs per MBMPO iter
+        "train_epochs": 500,
+        # Model Batch Size
+        "batch_size": 500,
+        # Training/Validation Split
+        "valid_split_ratio": 0.2,
+        # Normalize Data (obs, action, and deltas)
+        "normalize_data": True,
+    },
+    # Workers sample from dynamics models
+    "custom_vector_env": custom_model_vector_env,
+    # How many enviornments there are per worker (vectorized)
+    "num_worker_envs": 20,
+    # How many iterations through MAML per MBMPO iteration
+    "num_maml_steps": 10,
+})
+# __sphinx_doc_end__
+# yapf: enable
+
+# Select Metric Keys for MAML Stats Tracing
+METRICS_KEYS = [
+    "episode_reward_mean", "episode_reward_min", "episode_reward_max"
+]
+
+
+class MetaUpdate:
+    def __init__(self, workers, num_steps, maml_steps, metric_gen):
+        """Computes the MetaUpdate step in MAML, adapted for MBMPO
+        for multiple MAML Iterations
+
+        Arguments:
+            workers (WorkerSet): Set of Workers
+            num_steps (int): Number of meta-update steps per MAML Iteration
+            maml_steps (int): MAML Iterations per MBMPO Iteration
+            metric_gen (Iterator): Generates metrics dictionary
+
+        Returns:
+            metrics (dict): MBMPO metrics for logging
+        """
+        self.workers = workers
+        self.num_steps = num_steps
+        self.step_counter = 0
+        self.maml_optimizer_steps = maml_steps
+        self.metric_gen = metric_gen
+        self.metrics = {}
+
+    def __call__(self, data_tuple):
+        """Arguments:
+            data_tuple (tuple): 1st element is samples collected from MAML
+            Inner adaptation steps and 2nd element is accumulated metrics
+        """
+        # Metaupdate Step
+        print("Meta-Update Step")
+        samples = data_tuple[0]
+        adapt_metrics_dict = data_tuple[1]
+        self.postprocess_metrics(
+            adapt_metrics_dict, prefix="MAMLIter{}".format(self.step_counter))
+
+        # MAML Meta-update
+        for i in range(self.maml_optimizer_steps):
+            fetches = self.workers.local_worker().learn_on_batch(samples)
+        fetches = get_learner_stats(fetches)
+
+        # Update KLS
+        def update(pi, pi_id):
+            assert "inner_kl" not in fetches, (
+                "inner_kl should be nested under policy id key", fetches)
+            if pi_id in fetches:
+                assert "inner_kl" in fetches[pi_id], (fetches, pi_id)
+                pi.update_kls(fetches[pi_id]["inner_kl"])
+            else:
+                logger.warning("No data for {}, not updating kl".format(pi_id))
+
+        self.workers.local_worker().foreach_trainable_policy(update)
+
+        # Modify Reporting Metrics
+        metrics = _get_shared_metrics()
+        metrics.info[LEARNER_INFO] = fetches
+        metrics.counters[STEPS_TRAINED_COUNTER] += samples.count
+
+        if self.step_counter == self.num_steps:
+            td_metric = self.workers.local_worker().foreach_policy(
+                fit_dynamics)[0]
+
+            # Sync workers with meta policy
+            self.workers.sync_weights()
+
+            # Sync TD Models with workers
+            sync_ensemble(self.workers)
+            sync_stats(self.workers)
+
+            metrics.counters[STEPS_SAMPLED_COUNTER] = td_metric[
+                STEPS_SAMPLED_COUNTER]
+
+            # Modify to CollectMetrics
+            res = self.metric_gen.__call__(None)
+            res.update(self.metrics)
+            self.step_counter = 0
+            print("MB-MPO Iteration Completed")
+            return [res]
+        else:
+            print("MAML Iteration {} Completed".format(self.step_counter))
+            self.step_counter += 1
+
+            # Sync workers with meta policy
+            print("Syncing Weights with Workers")
+            self.workers.sync_weights()
+            return []
+
+    def postprocess_metrics(self, metrics, prefix=""):
+        """Appends prefix to current metrics
+
+        Arguments:
+            metrics (dict): Dictionary of current metrics
+            prefix (str): Prefix string to be appended
+        """
+        for key in metrics.keys():
+            self.metrics[prefix + "_" + key] = metrics[key]
+
+
+def post_process_metrics(prefix, workers, metrics):
+    """Update Current Dataset Metrics and filter out specific keys
+
+    Arguments:
+        prefix (str): Prefix string to be appended
+        workers (WorkerSet): Set of workers
+        metrics (dict): Current metrics dictionary
+    """
+    res = collect_metrics(remote_workers=workers.remote_workers())
+    for key in METRICS_KEYS:
+        metrics[prefix + "_" + key] = res[key]
+    return metrics
+
+
+def inner_adaptation(workers, samples):
+    # Each worker performs one gradient descent
+    for i, e in enumerate(workers.remote_workers()):
+        e.learn_on_batch.remote(samples[i])
+
+
+def fit_dynamics(policy, pid):
+    return policy.dynamics_model.fit()
+
+
+def sync_ensemble(workers):
+    """Syncs dynamics ensemble weights from main to workers
+
+    Arguments:
+        workers (WorkerSet): Set of workers, including main
+    """
+
+    def get_ensemble_weights(worker):
+        policy_map = worker.policy_map
+        policies = policy_map.keys()
+
+        def policy_ensemble_weights(policy):
+            model = policy.dynamics_model
+            return {
+                k: v.cpu().detach().numpy()
+                for k, v in model.state_dict().items()
+            }
+
+        return {
+            pid: policy_ensemble_weights(policy)
+            for pid, policy in policy_map.items() if pid in policies
+        }
+
+    def set_ensemble_weights(policy, pid, weights):
+        weights = weights[pid]
+        weights = convert_to_torch_tensor(weights, device=policy.device)
+        model = policy.dynamics_model
+        model.load_state_dict(weights)
+
+    if workers.remote_workers():
+        weights = ray.put(get_ensemble_weights(workers.local_worker()))
+        set_func = ray.put(set_ensemble_weights)
+        for e in workers.remote_workers():
+            e.foreach_policy.remote(set_func, weights=weights)
+
+
+def sync_stats(workers):
+    def get_normalizations(worker):
+        policy = worker.policy_map[DEFAULT_POLICY_ID]
+        return policy.dynamics_model.normalizations
+
+    def set_normalizations(policy, pid, normalizations):
+        policy.dynamics_model.set_norms(normalizations)
+
+    if workers.remote_workers():
+        normalization_dict = ray.put(
+            get_normalizations(workers.local_worker()))
+        set_func = ray.put(set_normalizations)
+        for e in workers.remote_workers():
+            e.foreach_policy.remote(
+                set_func, normalizations=normalization_dict)
+
+
+def post_process_samples(samples, config):
+    # Instead of using NN for value function, we use regression
+    split_lst = []
+    for sample in samples:
+        indexes = np.asarray(sample["dones"]).nonzero()[0]
+        indexes = indexes + 1
+
+        reward_list = np.split(sample["rewards"], indexes)[:-1]
+        observation_list = np.split(sample["obs"], indexes)[:-1]
+
+        paths = []
+        for i in range(0, len(reward_list)):
+            paths.append({
+                "rewards": reward_list[i],
+                "observations": observation_list[i]
+            })
+
+        paths = calculate_gae_advantages(paths, config["gamma"],
+                                         config["lambda"])
+
+        advantages = np.concatenate([path["advantages"] for path in paths])
+        sample["advantages"] = standardized(advantages)
+        split_lst.append(sample.count)
+    return samples, split_lst
+
+
+# Similar to MAML Execution Plan
+def execution_plan(workers, config):
+    # Train TD Models
+    workers.local_worker().foreach_policy(fit_dynamics)
+
+    # Sync workers policy with workers
+    workers.sync_weights()
+
+    # Sync TD Models and normalization stats with workers
+    sync_ensemble(workers)
+    sync_stats(workers)
+
+    # Dropping metrics from the first iteration
+    episodes, to_be_collected = collect_episodes(
+        workers.local_worker(),
+        workers.remote_workers(), [],
+        timeout_seconds=9999)
+
+    # Metrics Collector
+    metric_collect = CollectMetrics(
+        workers,
+        min_history=0,
+        timeout_seconds=config["collect_metrics_timeout"])
+
+    inner_steps = config["inner_adaptation_steps"]
+
+    def inner_adaptation_steps(itr):
+        buf = []
+        split = []
+        metrics = {}
+        for samples in itr:
+            print("Collecting Samples, Inner Adaptation {}".format(len(split)))
+            # Processing Samples (Standardize Advantages)
+            samples, split_lst = post_process_samples(samples, config)
+
+            buf.extend(samples)
+            split.append(split_lst)
+
+            adapt_iter = len(split) - 1
+            prefix = "DynaTrajInner_" + str(adapt_iter)
+            metrics = post_process_metrics(prefix, workers, metrics)
+
+            if len(split) > inner_steps:
+                out = SampleBatch.concat_samples(buf)
+                out["split"] = np.array(split)
+                buf = []
+                split = []
+
+                yield out, metrics
+                metrics = {}
+            else:
+                inner_adaptation(workers, samples)
+
+    # Iterator for Inner Adaptation Data gathering (from pre->post adaptation)
+    rollouts = from_actors(workers.remote_workers())
+    rollouts = rollouts.batch_across_shards()
+    rollouts = rollouts.transform(inner_adaptation_steps)
+
+    # Metaupdate Step with outer combine loop for multiple MAML iterations
+    train_op = rollouts.combine(
+        MetaUpdate(workers, config["num_maml_steps"],
+                   config["maml_optimizer_steps"], metric_collect))
+    return train_op
+
+
+def get_policy_class(config):
+    return MBMPOTorchPolicy
+
+
+def validate_config(config):
+    config["framework"] = "torch"
+    if config["framework"] != "torch":
+        raise ValueError("MB-MPO not supported in Tensorflow yet!")
+    if config["inner_adaptation_steps"] <= 0:
+        raise ValueError("Inner Adaptation Steps must be >=1.")
+    if config["maml_optimizer_steps"] <= 0:
+        raise ValueError("PPO steps for meta-update needs to be >=0")
+    if config["entropy_coeff"] < 0:
+        raise ValueError("entropy_coeff must be >=0")
+    if config["batch_mode"] != "complete_episodes":
+        raise ValueError("truncate_episodes not supported")
+    if config["num_workers"] <= 0:
+        raise ValueError("Must have at least 1 worker/task.")
+
+
+MBMPOTrainer = build_trainer(
+    name="MBMPO",
+    default_config=DEFAULT_CONFIG,
+    default_policy=MBMPOTorchPolicy,
+    get_policy_class=get_policy_class,
+    execution_plan=execution_plan,
+    validate_config=validate_config)

rllib/agents/mbmpo/mbmpo_torch_policy.py

@@ -0,0 +1,65 @@
+import logging
+
+import ray
+from ray.rllib.policy.torch_policy_template import build_torch_policy
+from ray.rllib.agents.ppo.ppo_tf_policy import postprocess_ppo_gae, \
+    setup_config
+from ray.rllib.agents.ppo.ppo_torch_policy import vf_preds_fetches
+from ray.rllib.agents.a3c.a3c_torch_policy import apply_grad_clipping
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.models.catalog import ModelCatalog
+from ray.rllib.agents.maml.maml_torch_policy import setup_mixins, \
+    maml_loss, maml_stats, maml_optimizer_fn, KLCoeffMixin
+
+torch, nn = try_import_torch()
+
+logger = logging.getLogger(__name__)
+
+
+def make_model_and_action_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")
+
+    # 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.
+    device = (torch.device("cuda")
+              if torch.cuda.is_available() else torch.device("cpu"))
+    policy.dynamics_model = ModelCatalog.get_model_v2(
+        obs_space,
+        action_space,
+        num_outputs=num_outputs,
+        model_config=config["dynamics_model"],
+        framework="torch",
+        name="dynamics_ensemble",
+    ).to(device)
+
+    action_dist, num_outputs = ModelCatalog.get_action_dist(
+        action_space, config, 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
+
+
+MBMPOTorchPolicy = build_torch_policy(
+    name="MBMPOTorchPolicy",
+    get_default_config=lambda: ray.rllib.agents.mbmpo.mbmpo.DEFAULT_CONFIG,
+    make_model_and_action_dist=make_model_and_action_dist,
+    loss_fn=maml_loss,
+    stats_fn=maml_stats,
+    optimizer_fn=maml_optimizer_fn,
+    extra_action_out_fn=vf_preds_fetches,
+    postprocess_fn=postprocess_ppo_gae,
+    extra_grad_process_fn=apply_grad_clipping,
+    before_init=setup_config,
+    after_init=setup_mixins,
+    mixins=[KLCoeffMixin])

rllib/agents/mbmpo/model_ensemble.py

@@ -0,0 +1,339 @@
+import gym
+from gym.spaces import Discrete, Box
+import numpy as np
+from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.evaluation.rollout_worker import get_global_worker
+from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.execution.common import STEPS_SAMPLED_COUNTER
+from ray.rllib.utils.types import SampleBatchType
+from ray.rllib.utils.torch_ops import convert_to_torch_tensor
+
+torch, nn = try_import_torch()
+
+
+class TDModel(nn.Module):
+    """Transition Dynamics Model (FC Network with Weight Norm)
+    """
+
+    def __init__(self,
+                 input_size,
+                 output_size,
+                 hidden_layers=[512, 512],
+                 hidden_nonlinearity=None,
+                 output_nonlinearity=None,
+                 weight_normalization=False,
+                 use_bias=True):
+
+        super().__init__()
+        assert len(hidden_layers) >= 1
+
+        if not hidden_nonlinearity:
+            hidden_nonlinearity = nn.ReLU
+
+        if weight_normalization:
+            weight_norm = nn.utils.weight_norm
+
+        self.layers = []
+        cur_size = input_size
+        for h_size in hidden_layers:
+            layer = nn.Linear(cur_size, h_size, bias=use_bias)
+            if weight_normalization:
+                layer = weight_norm(layer)
+            self.layers.append(layer)
+            if hidden_nonlinearity:
+                self.layers.append(hidden_nonlinearity())
+            cur_size = h_size
+
+        layer = nn.Linear(cur_size, output_size, bias=use_bias)
+        if weight_normalization:
+            layer = weight_norm(layer)
+        self.layers.append(layer)
+        if output_nonlinearity:
+            self.layers.append(output_nonlinearity())
+
+        self.model = nn.Sequential(*self.layers)
+
+    def forward(self, x):
+        return self.model(x)
+
+
+if torch:
+    class TDDataset(torch.utils.data.Dataset):
+        def __init__(self, dataset: SampleBatchType, norms):
+            self.count = dataset.count
+            obs = dataset[SampleBatch.CUR_OBS]
+            actions = dataset[SampleBatch.ACTIONS]
+            delta = dataset[SampleBatch.NEXT_OBS] - obs
+
+            if norms:
+                obs = normalize(obs, norms[SampleBatch.CUR_OBS])
+                actions = normalize(actions, norms[SampleBatch.ACTIONS])
+                delta = normalize(delta, norms["delta"])
+
+            self.x = np.concatenate([obs, actions], axis=1)
+            self.y = delta
+
+        def __len__(self):
+            return self.count
+
+        def __getitem__(self, index):
+            return self.x[index], self.y[index]
+
+
+def normalize(data_array, stats):
+    mean, std = stats
+    return (data_array - mean) / (std + 1e-10)
+
+
+def denormalize(data_array, stats):
+    mean, std = stats
+    return data_array * (std + 1e-10) + mean
+
+
+def mean_std_stats(dataset: SampleBatchType):
+    norm_dict = {}
+    obs = dataset[SampleBatch.CUR_OBS]
+    act = dataset[SampleBatch.ACTIONS]
+    delta = dataset[SampleBatch.NEXT_OBS] - obs
+
+    norm_dict[SampleBatch.CUR_OBS] = (np.mean(obs, axis=0), np.std(
+        obs, axis=0))
+    norm_dict[SampleBatch.ACTIONS] = (np.mean(act, axis=0), np.std(
+        act, axis=0))
+    norm_dict["delta"] = (np.mean(delta, axis=0), np.std(delta, axis=0))
+
+    return norm_dict
+
+
+def process_samples(samples: SampleBatchType):
+    filter_keys = [
+        SampleBatch.CUR_OBS, SampleBatch.ACTIONS, SampleBatch.NEXT_OBS
+    ]
+    filtered = {}
+    for key in filter_keys:
+        filtered[key] = samples[key]
+    return SampleBatch(filtered)
+
+
+class DynamicsEnsembleCustomModel(TorchModelV2, nn.Module):
+    """Represents a Transition Dyamics ensemble
+    """
+
+    def __init__(self, obs_space, action_space, num_outputs, model_config,
+                 name):
+        """Initializes a DynamicEnsemble object.
+        """
+        nn.Module.__init__(self)
+        if 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, ))
+        elif isinstance(action_space, Box):
+            input_space = gym.spaces.Box(
+                obs_space.low[0],
+                obs_space.high[0],
+                shape=(obs_space.shape[0] + action_space.shape[0], ))
+        super(DynamicsEnsembleCustomModel, self).__init__(
+            input_space, action_space, num_outputs, model_config, name)
+
+        self.num_models = model_config["ensemble_size"]
+        self.max_epochs = model_config["train_epochs"]
+        self.lr = model_config["lr"]
+        self.valid_split = model_config["valid_split_ratio"]
+        self.batch_size = model_config["batch_size"]
+        self.normalize_data = model_config["normalize_data"]
+        self.normalizations = {}
+        self.dynamics_ensemble = [
+            TDModel(
+                input_size=input_space.shape[0],
+                output_size=obs_space.shape[0],
+                hidden_layers=model_config["fcnet_hiddens"],
+                hidden_nonlinearity=nn.ReLU,
+                output_nonlinearity=None,
+                weight_normalization=True) for _ in range(self.num_models)
+        ]
+
+        for i in range(self.num_models):
+            self.add_module("TD-model-" + str(i), self.dynamics_ensemble[i])
+        self.replay_buffer_max = 100000
+        self.replay_buffer = None
+        self.optimizers = [
+            torch.optim.Adam(
+                self.dynamics_ensemble[i].parameters(), lr=self.lr)
+            for i in range(self.num_models)
+        ]
+        # Metric Reporting
+        self.metrics = {}
+        self.metrics[STEPS_SAMPLED_COUNTER] = 0
+
+        # For each worker, choose a random model to choose trajectories from
+        self.sample_index = np.random.randint(self.num_models)
+        self.global_itr = 0
+        self.device = (torch.device("cuda")
+                       if torch.cuda.is_available() else torch.device("cpu"))
+
+    def forward(self, x):
+        """Outputs the delta between next and current observation.
+        """
+        return self.dynamics_ensemble[self.sample_index](x)
+
+    # Loss functions for each TD model in Ensemble (Standard L2 Loss)
+    def loss(self, x, y):
+        xs = torch.chunk(x, self.num_models)
+        ys = torch.chunk(y, self.num_models)
+        return [
+            torch.mean(
+                torch.pow(self.dynamics_ensemble[i](xs[i]) - ys[i], 2.0))
+            for i in range(self.num_models)
+        ]
+
+    # Fitting Dynamics Ensembles per MBMPO Iter
+    def fit(self):
+        # Add env samples to Replay Buffer
+        local_worker = get_global_worker()
+        new_samples = local_worker.sample()
+        if not self.global_itr:
+            tmp = local_worker.sample()
+            new_samples.concat(tmp)
+
+        # Process Samples
+        new_samples = process_samples(new_samples)
+
+        if not self.replay_buffer:
+            self.replay_buffer = new_samples
+        else:
+            self.replay_buffer = self.replay_buffer.concat(new_samples)
+
+        # Keep Replay Buffer Size Constant
+        self.replay_buffer = self.replay_buffer.slice(
+            start=-self.replay_buffer_max, end=None)
+
+        if self.normalize_data:
+            self.normalizations = mean_std_stats(self.replay_buffer)
+
+        # Keep Track of Timesteps from Real Environment Timesteps Sampled
+        self.metrics[STEPS_SAMPLED_COUNTER] += new_samples.count
+
+        # Create Train and Val Datasets for each TD model
+        train_loaders = []
+        val_loaders = []
+        for i in range(self.num_models):
+            t, v = self.split_train_val(self.replay_buffer)
+            train_loaders.append(
+                torch.utils.data.DataLoader(
+                    TDDataset(t, self.normalizations),
+                    batch_size=self.batch_size,
+                    shuffle=True))
+            val_loaders.append(
+                torch.utils.data.DataLoader(
+                    TDDataset(v, self.normalizations),
+                    batch_size=v.count,
+                    shuffle=False))
+
+        # List of which models in ensemble to train
+        indexes = list(range(self.num_models))
+
+        valid_loss_roll_avg = None
+        roll_avg_persitency = 0.95
+
+        def convert_to_str(lst):
+            return " ".join([str(elem) for elem in lst])
+
+        for epoch in range(self.max_epochs):
+            # Training
+            for data in zip(*train_loaders):
+                x = torch.cat([d[0] for d in data], dim=0).to(self.device)
+                y = torch.cat([d[1] for d in data], dim=0).to(self.device)
+                train_losses = self.loss(x, y)
+                for ind in indexes:
+                    self.optimizers[ind].zero_grad()
+                    train_losses[ind].backward()
+                    self.optimizers[ind].step()
+
+                for ind in range(self.num_models):
+                    train_losses[ind] = train_losses[
+                        ind].detach().cpu().numpy()
+
+                del x
+                del y
+
+            # Validation
+            val_lists = []
+            for data in zip(*val_loaders):
+                x = torch.cat([d[0] for d in data], dim=0).to(self.device)
+                y = torch.cat([d[1] for d in data], dim=0).to(self.device)
+                val_losses = self.loss(x, y)
+                val_lists.append(val_losses)
+
+                for ind in indexes:
+                    self.optimizers[ind].zero_grad()
+
+                for ind in range(self.num_models):
+                    val_losses[ind] = val_losses[ind].detach().cpu().numpy()
+                del x
+                del y
+
+            val_lists = np.array(val_lists)
+            avg_val_losses = np.mean(val_lists, axis=0)
+
+            if valid_loss_roll_avg is None:
+                # Make sure that training doesnt end first epoch
+                valid_loss_roll_avg = 1.5 * avg_val_losses
+                valid_loss_roll_avg_prev = 2.0 * avg_val_losses
+
+            valid_loss_roll_avg = roll_avg_persitency*valid_loss_roll_avg + \
+                (1.0-roll_avg_persitency)*avg_val_losses
+
+            print("Training Dynamics Ensemble - Epoch #%i:"
+                  "Train loss: %s, Valid Loss: %s,  Moving Avg Valid Loss: %s"
+                  % (epoch, convert_to_str(train_losses),
+                     convert_to_str(avg_val_losses),
+                     convert_to_str(valid_loss_roll_avg)))
+            for i in range(self.num_models):
+                if (valid_loss_roll_avg_prev[i] < valid_loss_roll_avg[i]
+                        or epoch == self.max_epochs - 1) and i in indexes:
+                    indexes.remove(i)
+                    print("Stopping Training of Model %i" % i)
+            valid_loss_roll_avg_prev = valid_loss_roll_avg
+            if (len(indexes) == 0):
+                break
+
+        self.global_itr += 1
+        # Returns Metric Dictionary
+        return self.metrics
+
+    def split_train_val(self, samples: SampleBatchType):
+        dataset_size = samples.count
+        indices = np.arange(dataset_size)
+        np.random.shuffle(indices)
+        split_idx = int(dataset_size * (1 - self.valid_split))
+        idx_train = indices[:split_idx]
+        idx_test = indices[split_idx:]
+
+        train = {}
+        val = {}
+        for key in samples.keys():
+            train[key] = samples[key][idx_train, :]
+            val[key] = samples[key][idx_test, :]
+        return SampleBatch(train), SampleBatch(val)
+
+    """Used by worker who gather trajectories via TD models
+    """
+
+    def predict_model_batches(self, obs, actions, device=None):
+        pre_obs = obs
+        if self.normalize_data:
+            obs = normalize(obs, self.normalizations[SampleBatch.CUR_OBS])
+            actions = normalize(actions,
+                                self.normalizations[SampleBatch.ACTIONS])
+        x = np.concatenate([obs, actions], axis=-1)
+        x = convert_to_torch_tensor(x, device=device)
+        delta = self.forward(x).detach().numpy()
+        if self.normalize_data:
+            delta = denormalize(delta, self.normalizations["delta"])
+        return pre_obs + delta
+
+    def set_norms(self, normalization_dict):
+        self.normalizations = normalization_dict

rllib/agents/mbmpo/model_vector_env.py

@@ -0,0 +1,102 @@
+import logging
+import numpy as np
+
+from ray.rllib.utils.annotations import override
+from ray.rllib.env.vector_env import VectorEnv
+from ray.rllib.evaluation.rollout_worker import get_global_worker
+from ray.rllib.env.base_env import BaseEnv
+
+logger = logging.getLogger(__name__)
+
+
+def custom_model_vector_env(env):
+    """Returns a VectorizedEnv wrapper around the current envioronment
+    To obtain worker configs, one can call get_global_worker().
+    """
+    worker = get_global_worker()
+    worker_index = worker.worker_index
+    if worker_index:
+        env = _VectorizedModelGymEnv(
+            make_env=worker.make_env_fn,
+            existing_envs=[env],
+            num_envs=worker.num_envs,
+            observation_space=env.observation_space,
+            action_space=env.action_space,
+        )
+    return BaseEnv.to_base_env(
+        env,
+        make_env=worker.make_env_fn,
+        num_envs=worker.num_envs,
+        remote_envs=False,
+        remote_env_batch_wait_ms=0)
+
+
+class _VectorizedModelGymEnv(VectorEnv):
+    """Vectorized Environment Wrapper for MB-MPO. Primary change is
+    in the vector_step method, which calls the dynamics models for
+    """
+
+    def __init__(self,
+                 make_env=None,
+                 existing_envs=None,
+                 num_envs=1,
+                 *,
+                 observation_space=None,
+                 action_space=None,
+                 env_config=None):
+        self.make_env = make_env
+        self.envs = existing_envs
+        self.num_envs = num_envs
+        while len(self.envs) < num_envs:
+            self.envs.append(self.make_env(len(self.envs)))
+
+        super().__init__(
+            observation_space=observation_space
+            or self.envs[0].observation_space,
+            action_space=action_space or self.envs[0].action_space,
+            num_envs=num_envs)
+        worker = get_global_worker()
+        self.model, self.device = worker.foreach_policy(
+            lambda x, y: (x.dynamics_model, x.device))[0]
+
+    @override(VectorEnv)
+    def vector_reset(self):
+        self.cur_obs = [e.reset() for e in self.envs]
+        return self.cur_obs
+
+    @override(VectorEnv)
+    def reset_at(self, index):
+        obs = self.envs[index].reset()
+        self.cur_obs[index] = obs
+        return obs
+
+    @override(VectorEnv)
+    def vector_step(self, actions):
+        if self.cur_obs is None:
+            raise ValueError("Need to reset env first")
+
+        obs_batch = np.stack(self.cur_obs, axis=0)
+        action_batch = np.stack(actions, axis=0)
+
+        next_obs_batch = self.model.predict_model_batches(
+            obs_batch, action_batch, device=self.device)
+
+        next_obs_batch = np.clip(next_obs_batch, -50, 50)
+
+        rew_batch = self.envs[0].reward(obs_batch, action_batch,
+                                        next_obs_batch)
+
+        if hasattr(self.envs[0], "done"):
+            dones_batch = self.envs[0].done(next_obs_batch)
+        else:
+            dones_batch = np.asarray([False for _ in range(self.num_envs)])
+
+        info_batch = [{} for _ in range(self.num_envs)]
+
+        self.cur_obs = next_obs_batch
+
+        return list(obs_batch), list(rew_batch), list(dones_batch), info_batch
+
+    @override(VectorEnv)
+    def get_unwrapped(self):
+        return self.envs

rllib/agents/mbmpo/utils.py

@@ -0,0 +1,68 @@
+import numpy as np
+import scipy
+
+
+class LinearFeatureBaseline():
+    def __init__(self, reg_coeff=1e-5):
+        self._coeffs = None
+        self._reg_coeff = reg_coeff
+
+    def get_param_values(self, **tags):
+        return self._coeffs
+
+    def set_param_values(self, val, **tags):
+        self._coeffs = val
+
+    def _features(self, path):
+        o = np.clip(path["observations"], -10, 10)
+        ll = len(path["rewards"])
+        al = np.arange(ll).reshape(-1, 1) / 100.0
+        return np.concatenate(
+            [o, o**2, al, al**2, al**3,
+             np.ones((ll, 1))], axis=1)
+
+    def fit(self, paths):
+        featmat = np.concatenate([self._features(path) for path in paths])
+        returns = np.concatenate([path["returns"] for path in paths])
+        reg_coeff = self._reg_coeff
+        for _ in range(5):
+            self._coeffs = np.linalg.lstsq(
+                featmat.T.dot(featmat) +
+                reg_coeff * np.identity(featmat.shape[1]),
+                featmat.T.dot(returns))[0]
+            if not np.any(np.isnan(self._coeffs)):
+                break
+            reg_coeff *= 10
+
+    def predict(self, path):
+        if self._coeffs is None:
+            return np.zeros(len(path["rewards"]))
+        return self._features(path).dot(self._coeffs)
+
+
+def calculate_gae_advantages(paths, discount, gae_lambda):
+    baseline = LinearFeatureBaseline()
+
+    for idx, path in enumerate(paths):
+        path["returns"] = discount_cumsum(path["rewards"], discount)
+
+    baseline.fit(paths)
+    all_path_baselines = [baseline.predict(path) for path in paths]
+
+    for idx, path in enumerate(paths):
+        path_baselines = np.append(all_path_baselines[idx], 0)
+        deltas = path["rewards"] + \
+            discount * path_baselines[1:] - \
+            path_baselines[:-1]
+        path["advantages"] = discount_cumsum(deltas, discount * gae_lambda)
+    return paths
+
+
+def discount_cumsum(x, discount):
+    """
+        Returns:
+            (float) : y[t] - discount*y[t+1] = x[t] or rev(y)[t]
+            - discount*rev(y)[t-1] = rev(x)[t]
+        """
+    return scipy.signal.lfilter(
+        [1], [1, float(-discount)], x[::-1], axis=0)[::-1]

rllib/agents/registry.py

@@ -100,6 +100,11 @@ def _import_maml():
     return maml.MAMLTrainer
 
 
+def _import_mbmpo():
+    from ray.rllib.agents import mbmpo
+    return mbmpo.MBMPOTrainer
+
+
 ALGORITHMS = {
     "SAC": _import_sac,
     "DDPG": _import_ddpg,
@@ -120,6 +125,7 @@ ALGORITHMS = {
     "DDPPO": _import_ddppo,
     "MARWIL": _import_marwil,
     "MAML": _import_maml,
+    "MBMPO": _import_mbmpo,
 }
 
 

rllib/evaluation/rollout_worker.py

@@ -299,6 +299,7 @@ class RolloutWorker(ParallelIteratorWorker):
             enable_periodic_logging()
 
         env_context = EnvContext(env_config or {}, worker_index)
+        self.env_context = env_context
         self.policy_config: TrainerConfigDict = policy_config
         if callbacks:
             self.callbacks: "DefaultCallbacks" = callbacks()
@@ -362,7 +363,11 @@ class RolloutWorker(ParallelIteratorWorker):
             return wrap(
                 env_creator(
                     env_context.copy_with_overrides(
-                        vector_index=vector_index, remote=remote_worker_envs)))
+                        worker_index=worker_index,
+                        vector_index=vector_index,
+                        remote=remote_worker_envs)))
+
+        self.make_env_fn = make_env
 
         self.tf_sess = None
         policy_dict = _validate_and_canonicalize(policy, self.env)
@@ -405,14 +410,13 @@ class RolloutWorker(ParallelIteratorWorker):
             self.policy_map, self.preprocessors = self._build_policy_map(
                 policy_dict, policy_config)
 
-        if (ray.is_initialized() and
-                ray.worker._mode() != ray.worker.LOCAL_MODE):
+        if (ray.is_initialized()
+                and ray.worker._mode() != ray.worker.LOCAL_MODE):
             # Check available number of GPUs
             if not ray.get_gpu_ids():
-                logger.debug(
-                    "Creating policy evaluation worker {}".format(
-                        worker_index) +
-                    " on CPU (please ignore any CUDA init errors)")
+                logger.debug("Creating policy evaluation worker {}".format(
+                    worker_index) +
+                             " on CPU (please ignore any CUDA init errors)")
             elif (policy_config["framework"] in ["tf2", "tf", "tfe"] and
                   not tf.config.experimental.list_physical_devices("GPU")) or \
                     (policy_config["framework"] == "torch" and
@@ -421,9 +425,8 @@ class RolloutWorker(ParallelIteratorWorker):
                     "GPUs were assigned to this worker by Ray, but "
                     "your DL framework ({}) reports GPU acceleration is "
                     "disabled. This could be due to a bad CUDA- or {} "
-                    "installation.".format(
-                        policy_config["framework"],
-                        policy_config["framework"]))
+                    "installation.".format(policy_config["framework"],
+                                           policy_config["framework"]))
 
         self.multiagent: bool = set(
             self.policy_map.keys()) != {DEFAULT_POLICY_ID}
@@ -445,15 +448,20 @@ class RolloutWorker(ParallelIteratorWorker):
         if self.worker_index == 0:
             logger.info("Built filter map: {}".format(self.filters))
 
-        # Always use vector env for consistency even if num_envs = 1.
-        self.async_env: BaseEnv = BaseEnv.to_base_env(
-            self.env,
-            make_env=make_env,
-            num_envs=num_envs,
-            remote_envs=remote_worker_envs,
-            remote_env_batch_wait_ms=remote_env_batch_wait_ms)
         self.num_envs: int = num_envs
 
+        if "custom_vector_env" in policy_config:
+            custom_vec_wrapper = policy_config["custom_vector_env"]
+            self.async_env = custom_vec_wrapper(self.env)
+        else:
+            # Always use vector env for consistency even if num_envs = 1.
+            self.async_env: BaseEnv = BaseEnv.to_base_env(
+                self.env,
+                make_env=make_env,
+                num_envs=num_envs,
+                remote_envs=remote_worker_envs,
+                remote_env_batch_wait_ms=remote_env_batch_wait_ms)
+
         # `truncate_episodes`: Allow a batch to contain more than one episode
         # (fragments) and always make the batch `rollout_fragment_length`
         # long.
@@ -821,20 +829,25 @@ class RolloutWorker(ParallelIteratorWorker):
     @DeveloperAPI
     def for_policy(self,
                    func: Callable[[Policy], T],
-                   policy_id: Optional[PolicyID] = DEFAULT_POLICY_ID) -> T:
+                   policy_id: Optional[PolicyID] = DEFAULT_POLICY_ID,
+                   **kwargs) -> T:
         """Apply the given function to the specified policy."""
 
-        return func(self.policy_map[policy_id])
+        return func(self.policy_map[policy_id], **kwargs)
 
     @DeveloperAPI
-    def foreach_policy(self, func: Callable[[Policy, PolicyID], T]) -> List[T]:
+    def foreach_policy(self, func: Callable[[Policy, PolicyID], T],
+                       **kwargs) -> List[T]:
         """Apply the given function to each (policy, policy_id) tuple."""
 
-        return [func(policy, pid) for pid, policy in self.policy_map.items()]
+        return [
+            func(policy, pid, **kwargs)
+            for pid, policy in self.policy_map.items()
+        ]
 
     @DeveloperAPI
-    def foreach_trainable_policy(
-            self, func: Callable[[Policy, PolicyID], T]) -> List[T]:
+    def foreach_trainable_policy(self, func: Callable[[Policy, PolicyID], T],
+                                 **kwargs) -> List[T]:
         """
         Applies the given function to each (policy, policy_id) tuple, which
         can be found in `self.policies_to_train`.
@@ -848,7 +861,8 @@ class RolloutWorker(ParallelIteratorWorker):
                 `func([policy], [ID])`-calls.
         """
         return [
-            func(policy, pid) for pid, policy in self.policy_map.items()
+            func(policy, pid, **kwargs)
+            for pid, policy in self.policy_map.items()
             if pid in self.policies_to_train
         ]
 

rllib/examples/env/halfcheetah.py

@@ -0,0 +1,49 @@
+import numpy as np
+from gym.envs.mujoco import HalfCheetahEnv
+import inspect
+
+
+def get_all_function_arguments(function, locals):
+    kwargs_dict = {}
+    for arg in inspect.getfullargspec(function).kwonlyargs:
+        if arg not in ["args", "kwargs"]:
+            kwargs_dict[arg] = locals[arg]
+    args = [locals[arg] for arg in inspect.getfullargspec(function).args]
+
+    if "args" in locals:
+        args += locals["args"]
+
+    if "kwargs" in locals:
+        kwargs_dict.update(locals["kwargs"])
+    return args, kwargs_dict
+
+
+class HalfCheetahWrapper(HalfCheetahEnv):
+    """HalfCheetah Wrapper that wraps Mujoco Halfcheetah-v2 env
+    with an additional defined reward function for model-based RL.
+
+    This is currently used for MBMPO.
+    """
+
+    def __init__(self, *args, **kwargs):
+        HalfCheetahEnv.__init__(self, *args, **kwargs)
+
+    def reward(self, obs, action, obs_next):
+        if obs.ndim == 2 and action.ndim == 2:
+            assert obs.shape == obs_next.shape
+            forward_vel = obs_next[:, 8]
+            ctrl_cost = 0.1 * np.sum(np.square(action), axis=1)
+            reward = forward_vel - ctrl_cost
+            return np.minimum(np.maximum(-1000.0, reward), 1000.0)
+        else:
+            forward_vel = obs_next[8]
+            ctrl_cost = 0.1 * np.square(action).sum()
+            reward = forward_vel - ctrl_cost
+            return np.minimum(np.maximum(-1000.0, reward), 1000.0)
+
+
+if __name__ == "__main__":
+    env = HalfCheetahWrapper()
+    env.reset()
+    for _ in range(1000):
+        env.step(env.action_space.sample())

rllib/tuned_examples/mbmpo/halfcheetah-mbmpo.yaml

@@ -0,0 +1,25 @@
+halfcheetah-mb-mpo:
+    env: ray.rllib.examples.env.halfcheetah.HalfCheetahWrapper
+    run: MBMPO
+    stop:
+        training_iteration: 500
+    config:
+         # Only supported in torch right now
+        framework: torch
+        horizon: 200
+        num_envs_per_worker: 20
+        inner_adaptation_steps: 1
+        maml_optimizer_steps: 8
+        gamma: 0.99
+        lambda: 1.0
+        lr: 0.001
+        clip_param: 0.5
+        kl_target: 0.01
+        kl_coeff: 0.0000000001
+        num_workers: 20
+        num_gpus: 1
+        inner_lr: 0.001
+        clip_actions: False
+        model:
+            fcnet_hiddens: [32, 32]
+            free_log_std: True

rllib/utils/exploration/gaussian_noise.py

@@ -120,7 +120,7 @@ class GaussianNoise(Exploration):
             true_fn=lambda: stochastic_actions,
             false_fn=lambda: deterministic_actions)
         # Logp=always zero.
-        logp = tf.zeros(shape=(batch_size,), dtype=tf.float32)
+        logp = tf.zeros(shape=(batch_size, ), dtype=tf.float32)
 
         # Increment `last_timestep` by 1 (or set to `timestep`).
         if self.framework in ["tf2", "tfe"]:
@@ -130,9 +130,9 @@ class GaussianNoise(Exploration):
                 self.last_timestep.assign(timestep)
             return action, logp
         else:
-            assign_op = (
-                tf1.assign_add(self.last_timestep, 1) if timestep is None else
-                tf1.assign(self.last_timestep, timestep))
+            assign_op = (tf1.assign_add(self.last_timestep, 1)
+                         if timestep is None else tf1.assign(
+                             self.last_timestep, timestep))
             with tf1.control_dependencies([assign_op]):
                 return action, logp
 

rllib/utils/exploration/ornstein_uhlenbeck_noise.py

@@ -126,7 +126,7 @@ class OrnsteinUhlenbeckNoise(GaussianNoise):
             false_fn=lambda: deterministic_actions)
         # Logp=always zero.
         batch_size = tf.shape(deterministic_actions)[0]
-        logp = tf.zeros(shape=(batch_size,), dtype=tf.float32)
+        logp = tf.zeros(shape=(batch_size, ), dtype=tf.float32)
 
         # Increment `last_timestep` by 1 (or set to `timestep`).
         if self.framework in ["tf2", "tfe"]:
@@ -136,9 +136,9 @@ class OrnsteinUhlenbeckNoise(GaussianNoise):
                 self.last_timestep = timestep
             return action, logp
         else:
-            assign_op = (
-                tf1.assign_add(self.last_timestep, 1) if timestep is None else
-                tf1.assign(self.last_timestep, timestep))
+            assign_op = (tf1.assign_add(self.last_timestep, 1)
+                         if timestep is None else tf1.assign(
+                             self.last_timestep, timestep))
             with tf1.control_dependencies([assign_op, ou_state_new]):
                 return action, logp