diff --git a/rllib/agents/maml/maml_torch_policy.py b/rllib/agents/maml/maml_torch_policy.py index e46876a5c..cf378a4ba 100644 --- a/rllib/agents/maml/maml_torch_policy.py +++ b/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 diff --git a/rllib/agents/mbmpo/__init__.py b/rllib/agents/mbmpo/__init__.py new file mode 100644 index 000000000..00c228fe5 --- /dev/null +++ b/rllib/agents/mbmpo/__init__.py @@ -0,0 +1,6 @@ +from ray.rllib.agents.mbmpo.mbmpo import MBMPOTrainer, DEFAULT_CONFIG + +__all__ = [ + "MBMPOTrainer", + "DEFAULT_CONFIG", +] diff --git a/rllib/agents/mbmpo/mbmpo.py b/rllib/agents/mbmpo/mbmpo.py new file mode 100644 index 000000000..cf24f8a78 --- /dev/null +++ b/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) diff --git a/rllib/agents/mbmpo/mbmpo_torch_policy.py b/rllib/agents/mbmpo/mbmpo_torch_policy.py new file mode 100644 index 000000000..5e0e6e205 --- /dev/null +++ b/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]) diff --git a/rllib/agents/mbmpo/model_ensemble.py b/rllib/agents/mbmpo/model_ensemble.py new file mode 100644 index 000000000..7caada9dd --- /dev/null +++ b/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 diff --git a/rllib/agents/mbmpo/model_vector_env.py b/rllib/agents/mbmpo/model_vector_env.py new file mode 100644 index 000000000..655169e06 --- /dev/null +++ b/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 diff --git a/rllib/agents/mbmpo/utils.py b/rllib/agents/mbmpo/utils.py new file mode 100644 index 000000000..16bb922da --- /dev/null +++ b/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] diff --git a/rllib/agents/registry.py b/rllib/agents/registry.py index 98be4a61b..a1dbe2b2d 100644 --- a/rllib/agents/registry.py +++ b/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, } diff --git a/rllib/evaluation/rollout_worker.py b/rllib/evaluation/rollout_worker.py index 021a8041f..7276a0ec5 100644 --- a/rllib/evaluation/rollout_worker.py +++ b/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 ] diff --git a/rllib/examples/env/halfcheetah.py b/rllib/examples/env/halfcheetah.py new file mode 100644 index 000000000..70f946468 --- /dev/null +++ b/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()) diff --git a/rllib/tuned_examples/mbmpo/halfcheetah-mbmpo.yaml b/rllib/tuned_examples/mbmpo/halfcheetah-mbmpo.yaml new file mode 100644 index 000000000..9e69fde03 --- /dev/null +++ b/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 diff --git a/rllib/utils/exploration/gaussian_noise.py b/rllib/utils/exploration/gaussian_noise.py index a3e5d32ab..53046348a 100644 --- a/rllib/utils/exploration/gaussian_noise.py +++ b/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 diff --git a/rllib/utils/exploration/ornstein_uhlenbeck_noise.py b/rllib/utils/exploration/ornstein_uhlenbeck_noise.py index 72e4c09b7..661057999 100644 --- a/rllib/utils/exploration/ornstein_uhlenbeck_noise.py +++ b/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