From 42cd414e5bf5a224703e8809881a035d5996f861 Mon Sep 17 00:00:00 2001 From: Michael Luo Date: Wed, 30 Dec 2020 07:11:57 -0800 Subject: [PATCH] [RLlib] New Offline RL Algorithm: CQL (based on SAC) (#13118) --- rllib/agents/cql/__init__.py | 8 + rllib/agents/cql/cql.py | 51 +++ rllib/agents/cql/cql_torch_policy.py | 301 ++++++++++++++++++ rllib/agents/registry.py | 6 + rllib/agents/sac/sac_torch_policy.py | 12 +- rllib/tuned_examples/cql/halfcheetah-cql.yaml | 45 +++ 6 files changed, 417 insertions(+), 6 deletions(-) create mode 100644 rllib/agents/cql/__init__.py create mode 100644 rllib/agents/cql/cql.py create mode 100644 rllib/agents/cql/cql_torch_policy.py create mode 100644 rllib/tuned_examples/cql/halfcheetah-cql.yaml diff --git a/rllib/agents/cql/__init__.py b/rllib/agents/cql/__init__.py new file mode 100644 index 000000000..154804137 --- /dev/null +++ b/rllib/agents/cql/__init__.py @@ -0,0 +1,8 @@ +from ray.rllib.agents.cql.cql import CQLTrainer, CQL_DEFAULT_CONFIG +from ray.rllib.agents.cql.cql_torch_policy import CQLTorchPolicy + +__all__ = [ + "CQL_DEFAULT_CONFIG", + "CQLTorchPolicy", + "CQLTrainer", +] diff --git a/rllib/agents/cql/cql.py b/rllib/agents/cql/cql.py new file mode 100644 index 000000000..04a63be72 --- /dev/null +++ b/rllib/agents/cql/cql.py @@ -0,0 +1,51 @@ +"""CQL (derived from SAC). +""" +from typing import Optional, Type + +from ray.rllib.agents.sac.sac import SACTrainer, \ + DEFAULT_CONFIG as SAC_CONFIG +from ray.rllib.agents.cql.cql_torch_policy import CQLTorchPolicy +from ray.rllib.utils.typing import TrainerConfigDict +from ray.rllib.policy.policy import Policy +from ray.rllib.utils import merge_dicts + +# yapf: disable +# __sphinx_doc_begin__ +CQL_DEFAULT_CONFIG = merge_dicts( + SAC_CONFIG, { + # You should override this to point to an offline dataset. + "input": "sampler", + # Number of iterations with Behavior Cloning Pretraining + "bc_iters": 20000, + # CQL Loss Temperature + "temperature": 1.0, + # Num Actions to sample for CQL Loss + "num_actions": 10, + # Whether to use the Langrangian for Alpha Prime (in CQL Loss) + "lagrangian": False, + # Lagrangian Threshold + "lagrangian_thresh": 5.0, + # Min Q Weight multiplier + "min_q_weight": 5.0, + }) +# __sphinx_doc_end__ +# yapf: enable + + +def validate_config(config: TrainerConfigDict): + if config["framework"] == "tf": + raise ValueError("Tensorflow CQL not implemented yet!") + + +def get_policy_class(config: TrainerConfigDict) -> Optional[Type[Policy]]: + if config["framework"] == "torch": + return CQLTorchPolicy + + +CQLTrainer = SACTrainer.with_updates( + name="CQL", + default_config=CQL_DEFAULT_CONFIG, + validate_config=validate_config, + default_policy=CQLTorchPolicy, + get_policy_class=get_policy_class, +) diff --git a/rllib/agents/cql/cql_torch_policy.py b/rllib/agents/cql/cql_torch_policy.py new file mode 100644 index 000000000..14812230f --- /dev/null +++ b/rllib/agents/cql/cql_torch_policy.py @@ -0,0 +1,301 @@ +""" +PyTorch policy class used for CQL. +""" +import numpy as np +import gym +import logging +from typing import Dict, List, Tuple, Type, Union + +import ray +import ray.experimental.tf_utils +from ray.rllib.agents.a3c.a3c_torch_policy import apply_grad_clipping +from ray.rllib.agents.sac.sac_tf_policy import postprocess_trajectory, \ + validate_spaces +from ray.rllib.agents.sac.sac_torch_policy import _get_dist_class, stats, \ + build_sac_model_and_action_dist, optimizer_fn, ComputeTDErrorMixin, \ + TargetNetworkMixin, setup_late_mixins, action_distribution_fn +from ray.rllib.models.torch.torch_action_dist import TorchDistributionWrapper +from ray.rllib.policy.policy_template import build_policy_class +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.policy.policy import Policy +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.typing import LocalOptimizer, TensorType, \ + TrainerConfigDict +from ray.rllib.utils.torch_ops import convert_to_torch_tensor + +torch, nn = try_import_torch() +F = nn.functional + +logger = logging.getLogger(__name__) + + +# Returns policy tiled actions and log probabilities for CQL Loss +def policy_actions_repeat(model, action_dist, obs, num_repeat=1): + obs_temp = obs.unsqueeze(1).repeat(1, num_repeat, 1).view( + obs.shape[0] * num_repeat, obs.shape[1]) + policy_dist = action_dist(model.get_policy_output(obs_temp), model) + actions = policy_dist.sample() + log_p = torch.unsqueeze(policy_dist.logp(actions), -1) + return actions, log_p.squeeze() + + +def q_values_repeat(model, obs, actions, twin=False): + action_shape = actions.shape[0] + obs_shape = obs.shape[0] + num_repeat = int(action_shape / obs_shape) + obs_temp = obs.unsqueeze(1).repeat(1, num_repeat, 1).view( + obs.shape[0] * num_repeat, obs.shape[1]) + if twin: + preds = model.get_q_values(obs_temp, actions) + else: + preds = model.get_twin_q_values(obs_temp, actions) + preds = preds.view(obs.shape[0], num_repeat, 1) + return preds + + +def cql_loss(policy: Policy, model: ModelV2, + dist_class: Type[TorchDistributionWrapper], + train_batch: SampleBatch) -> Union[TensorType, List[TensorType]]: + print(policy.cur_iter) + policy.cur_iter += 1 + # For best performance, turn deterministic off + deterministic = policy.config["_deterministic_loss"] + twin_q = policy.config["twin_q"] + discount = policy.config["gamma"] + action_low = model.action_space.low[0] + action_high = model.action_space.high[0] + + # CQL Parameters + bc_iters = policy.config["bc_iters"] + cql_temp = policy.config["temperature"] + num_actions = policy.config["num_actions"] + min_q_weight = policy.config["min_q_weight"] + use_lagrange = policy.config["lagrangian"] + target_action_gap = policy.config["lagrangian_thresh"] + + obs = train_batch[SampleBatch.CUR_OBS] + actions = train_batch[SampleBatch.ACTIONS] + rewards = train_batch[SampleBatch.REWARDS] + next_obs = train_batch[SampleBatch.NEXT_OBS] + terminals = train_batch[SampleBatch.DONES] + + model_out_t, _ = model({ + "obs": obs, + "is_training": True, + }, [], None) + + model_out_tp1, _ = model({ + "obs": next_obs, + "is_training": True, + }, [], None) + + target_model_out_tp1, _ = policy.target_model({ + "obs": next_obs, + "is_training": True, + }, [], None) + + action_dist_class = _get_dist_class(policy.config, policy.action_space) + action_dist_t = action_dist_class( + model.get_policy_output(model_out_t), policy.model) + policy_t = action_dist_t.sample() if not deterministic else \ + action_dist_t.deterministic_sample() + log_pis_t = torch.unsqueeze(action_dist_t.logp(policy_t), -1) + + # Unlike original SAC, Alpha and Actor Loss are computed first. + # Alpha Loss + alpha_loss = -(model.log_alpha * + (log_pis_t + model.target_entropy).detach()).mean() + + # Policy Loss (Either Behavior Clone Loss or SAC Loss) + alpha = torch.exp(model.log_alpha) + if policy.cur_iter >= bc_iters: + min_q = model.get_q_values(model_out_t, policy_t) + if twin_q: + twin_q = model.get_twin_q_values(model_out_t, policy_t) + min_q = torch.min(min_q, twin_q) + actor_loss = (alpha.detach() * log_pis_t - min_q).mean() + else: + bc_logp = action_dist_t.logp(actions) + actor_loss = (alpha * log_pis_t - bc_logp).mean() + + # Critic Loss (Standard SAC Critic L2 Loss + CQL Entropy Loss) + # SAC Loss + action_dist_tp1 = action_dist_class( + model.get_policy_output(model_out_tp1), policy.model) + policy_tp1 = action_dist_tp1.sample() if not deterministic else \ + action_dist_tp1.deterministic_sample() + + # Q-values for the batched actions. + q_t = model.get_q_values(model_out_t, train_batch[SampleBatch.ACTIONS]) + q_t = torch.squeeze(q_t, dim=-1) + if twin_q: + twin_q_t = model.get_twin_q_values(model_out_t, + train_batch[SampleBatch.ACTIONS]) + twin_q_t = torch.squeeze(twin_q_t, dim=-1) + + # Target q network evaluation. + q_tp1 = policy.target_model.get_q_values(target_model_out_tp1, policy_tp1) + if twin_q: + twin_q_tp1 = policy.target_model.get_twin_q_values( + target_model_out_tp1, policy_tp1) + # Take min over both twin-NNs. + q_tp1 = torch.min(q_tp1, twin_q_tp1) + q_tp1 = torch.squeeze(input=q_tp1, dim=-1) + q_tp1 = (1.0 - terminals.float()) * q_tp1 + + # compute RHS of bellman equation + q_t_target = ( + rewards + (discount**policy.config["n_step"]) * q_tp1).detach() + + # Compute the TD-error (potentially clipped), for priority replay buffer + base_td_error = torch.abs(q_t - q_t_target) + if twin_q: + twin_td_error = torch.abs(twin_q_t - q_t_target) + td_error = 0.5 * (base_td_error + twin_td_error) + else: + td_error = base_td_error + critic_loss = [nn.MSELoss()(q_t, q_t_target)] + if twin_q: + critic_loss.append(nn.MSELoss()(twin_q_t, q_t_target)) + + # CQL Loss (We are using Entropy version of CQL (the best version)) + rand_actions = convert_to_torch_tensor( + torch.FloatTensor(actions.shape[0] * num_actions, + actions.shape[-1]).uniform_(action_low, action_high)) + curr_actions, curr_logp = policy_actions_repeat(model, action_dist_class, + obs, num_actions) + next_actions, next_logp = policy_actions_repeat(model, action_dist_class, + next_obs, num_actions) + curr_logp = curr_logp.view(actions.shape[0], num_actions, 1) + next_logp = next_logp.view(actions.shape[0], num_actions, 1) + + q1_rand = q_values_repeat(model, model_out_t, rand_actions) + q1_curr_actions = q_values_repeat(model, model_out_t, curr_actions) + q1_next_actions = q_values_repeat(model, model_out_t, next_actions) + + if twin_q: + q2_rand = q_values_repeat(model, model_out_t, rand_actions, twin=True) + q2_curr_actions = q_values_repeat( + model, model_out_t, curr_actions, twin=True) + q2_next_actions = q_values_repeat( + model, model_out_t, next_actions, twin=True) + + random_density = np.log(0.5**curr_actions.shape[-1]) + cat_q1 = torch.cat([ + q1_rand - random_density, q1_next_actions - next_logp.detach(), + q1_curr_actions - curr_logp.detach() + ], 1) + if twin_q: + cat_q2 = torch.cat([ + q2_rand - random_density, q2_next_actions - next_logp.detach(), + q2_curr_actions - curr_logp.detach() + ], 1) + + min_qf1_loss = torch.logsumexp( + cat_q1 / cql_temp, dim=1).mean() * min_q_weight * cql_temp + min_qf1_loss = min_qf1_loss - q_t.mean() * min_q_weight + if twin_q: + min_qf2_loss = torch.logsumexp( + cat_q2 / cql_temp, dim=1).mean() * min_q_weight * cql_temp + min_qf2_loss = min_qf2_loss - twin_q_t.mean() * min_q_weight + + if use_lagrange: + alpha_prime = torch.clamp( + model.log_alpha_prime.exp(), min=0.0, max=1000000.0)[0] + min_qf1_loss = alpha_prime * (min_qf1_loss - target_action_gap) + if twin_q: + min_qf2_loss = alpha_prime * (min_qf2_loss - target_action_gap) + alpha_prime_loss = 0.5 * (-min_qf1_loss - min_qf2_loss) + else: + alpha_prime_loss = -min_qf1_loss + + cql_loss = [min_qf2_loss] + if twin_q: + cql_loss.append(min_qf2_loss) + + critic_loss[0] += min_qf1_loss + if twin_q: + critic_loss[1] += min_qf2_loss + + # Save for stats function. + policy.q_t = q_t + policy.policy_t = policy_t + policy.log_pis_t = log_pis_t + policy.td_error = td_error + policy.actor_loss = actor_loss + policy.critic_loss = critic_loss + policy.alpha_loss = alpha_loss + policy.log_alpha_value = model.log_alpha + policy.alpha_value = alpha + policy.target_entropy = model.target_entropy + # CQL Stats + policy.cql_loss = cql_loss + if use_lagrange: + policy.log_alpha_prime_value = model.log_alpha_prime[0] + policy.alpha_prime_value = alpha_prime + policy.alpha_prime_loss = alpha_prime_loss + + # Return all loss terms corresponding to our optimizers. + if use_lagrange: + return tuple([policy.actor_loss] + policy.critic_loss + + [policy.alpha_loss] + [policy.alpha_prime_loss]) + return tuple([policy.actor_loss] + policy.critic_loss + + [policy.alpha_loss]) + + +def cql_stats(policy: Policy, + train_batch: SampleBatch) -> Dict[str, TensorType]: + sac_dict = stats(policy, train_batch) + sac_dict["cql_loss"] = torch.mean(torch.stack(policy.cql_loss)) + if policy.config["lagrangian"]: + sac_dict["log_alpha_prime_value"] = policy.log_alpha_prime_value + sac_dict["alpha_prime_value"] = policy.alpha_prime_value + sac_dict["alpha_prime_loss"] = policy.alpha_prime_loss + return sac_dict + + +def cql_optimizer_fn(policy: Policy, config: TrainerConfigDict) -> \ + Tuple[LocalOptimizer]: + policy.cur_iter = 0 + opt_list = optimizer_fn(policy, config) + if config["lagrangian"]: + log_alpha_prime = nn.Parameter( + torch.zeros(1, requires_grad=True).float()) + policy.model.register_parameter("log_alpha_prime", log_alpha_prime) + policy.alpha_prime_optim = torch.optim.Adam( + params=[policy.model.log_alpha_prime], + lr=config["optimization"]["critic_learning_rate"], + eps=1e-7, # to match tf.keras.optimizers.Adam's epsilon default + ) + return tuple([policy.actor_optim] + policy.critic_optims + + [policy.alpha_optim] + [policy.alpha_prime_optim]) + return opt_list + + +def cql_setup_late_mixins(policy: Policy, obs_space: gym.spaces.Space, + action_space: gym.spaces.Space, + config: TrainerConfigDict) -> None: + setup_late_mixins(policy, obs_space, action_space, config) + if config["lagrangian"]: + policy.model.log_alpha_prime = policy.model.log_alpha_prime.to( + policy.device) + + +# Build a child class of `TorchPolicy`, given the custom functions defined +# above. +CQLTorchPolicy = build_policy_class( + name="CQLTorchPolicy", + framework="torch", + loss_fn=cql_loss, + get_default_config=lambda: ray.rllib.agents.cql.cql.CQL_DEFAULT_CONFIG, + stats_fn=cql_stats, + postprocess_fn=postprocess_trajectory, + extra_grad_process_fn=apply_grad_clipping, + optimizer_fn=cql_optimizer_fn, + validate_spaces=validate_spaces, + before_loss_init=cql_setup_late_mixins, + make_model_and_action_dist=build_sac_model_and_action_dist, + mixins=[TargetNetworkMixin, ComputeTDErrorMixin], + action_distribution_fn=action_distribution_fn, +) diff --git a/rllib/agents/registry.py b/rllib/agents/registry.py index 924711cf0..8ec4a4582 100644 --- a/rllib/agents/registry.py +++ b/rllib/agents/registry.py @@ -40,6 +40,11 @@ def _import_bc(): return marwil.BCTrainer +def _import_cql(): + from ray.rllib.agents import cql + return cql.CQLTrainer + + def _import_ddpg(): from ray.rllib.agents import ddpg return ddpg.DDPGTrainer @@ -128,6 +133,7 @@ ALGORITHMS = { "APPO": _import_appo, "ARS": _import_ars, "BC": _import_bc, + "CQL": _import_cql, "ES": _import_es, "DDPG": _import_ddpg, "DDPPO": _import_ddppo, diff --git a/rllib/agents/sac/sac_torch_policy.py b/rllib/agents/sac/sac_torch_policy.py index d1d53697b..49bc65557 100644 --- a/rllib/agents/sac/sac_torch_policy.py +++ b/rllib/agents/sac/sac_torch_policy.py @@ -441,12 +441,12 @@ class TargetNetworkMixin: model_state_dict = self.model.state_dict() # Support partial (soft) synching. # If tau == 1.0: Full sync from Q-model to target Q-model. - if tau != 1.0: - target_state_dict = self.target_model.state_dict() - model_state_dict = { - k: tau * model_state_dict[k] + (1 - tau) * v - for k, v in target_state_dict.items() - } + target_state_dict = self.target_model.state_dict() + model_state_dict = { + k: tau * model_state_dict[k] + (1 - tau) * v + for k, v in target_state_dict.items() + } + self.target_model.load_state_dict(model_state_dict) diff --git a/rllib/tuned_examples/cql/halfcheetah-cql.yaml b/rllib/tuned_examples/cql/halfcheetah-cql.yaml new file mode 100644 index 000000000..5bab20751 --- /dev/null +++ b/rllib/tuned_examples/cql/halfcheetah-cql.yaml @@ -0,0 +1,45 @@ +halfcheetah_cql: + env: HalfCheetah-v3 + run: CQL + stop: + episode_reward_mean: 9000 + config: + # SAC Configs + framework: torch + horizon: 1000 + soft_horizon: false + Q_model: + fcnet_activation: relu + fcnet_hiddens: [256, 256] + policy_model: + fcnet_activation: relu + fcnet_hiddens: [256, 256] + tau: 0.005 + target_entropy: auto + no_done_at_end: false + n_step: 1 + rollout_fragment_length: 1 + prioritized_replay: true + train_batch_size: 256 + target_network_update_freq: 1 + timesteps_per_iteration: 1000 + learning_starts: 10000 + optimization: + actor_learning_rate: 0.0003 + critic_learning_rate: 0.0003 + entropy_learning_rate: 0.0003 + num_workers: 0 + num_gpus: 0 + clip_actions: false + normalize_actions: true + evaluation_interval: 1 + metrics_smoothing_episodes: 5 + # CQL Configs + min_q_weight: 5.0 + bc_iters: 20000 + temperature: 1.0 + num_actions: 10 + lagrangian: False + evaluation_config: + input: sampler +