chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,13 @@
|
||||
# Conservative Q-Learning (CQL)
|
||||
|
||||
## Overview
|
||||
|
||||
[CQL](https://arxiv.org/abs/2006.04779) is an offline RL algorithm that mitigates the overestimation of Q-values outside the dataset distribution via convservative critic estimates. CQL does this by adding a simple Q regularizer loss to the standard Belman update loss. This ensures that the critic does not output overly-optimistic Q-values and can be added on top of any off-policy Q-learning algorithm (in this case, we use SAC).
|
||||
|
||||
## Documentation & Implementation:
|
||||
|
||||
Conservative Q-Learning (CQL).
|
||||
|
||||
**[Detailed Documentation](https://docs.ray.io/en/master/rllib-algorithms.html#cql)**
|
||||
|
||||
**[Implementation](https://github.com/ray-project/ray/blob/master/rllib/algorithms/cql/cql.py)**
|
||||
@@ -0,0 +1,9 @@
|
||||
from ray.rllib.algorithms.cql.cql import CQL, CQLConfig
|
||||
from ray.rllib.algorithms.cql.cql_torch_policy import CQLTorchPolicy
|
||||
|
||||
__all__ = [
|
||||
"CQL",
|
||||
"CQLConfig",
|
||||
# @OldAPIStack
|
||||
"CQLTorchPolicy",
|
||||
]
|
||||
@@ -0,0 +1,378 @@
|
||||
import logging
|
||||
from typing import Optional, Type, Union
|
||||
|
||||
from typing_extensions import Self
|
||||
|
||||
from ray._common.deprecation import (
|
||||
DEPRECATED_VALUE,
|
||||
deprecation_warning,
|
||||
)
|
||||
from ray.rllib.algorithms.algorithm_config import AlgorithmConfig, NotProvided
|
||||
from ray.rllib.algorithms.cql.cql_tf_policy import CQLTFPolicy
|
||||
from ray.rllib.algorithms.cql.cql_torch_policy import CQLTorchPolicy
|
||||
from ray.rllib.algorithms.sac.sac import (
|
||||
SAC,
|
||||
SACConfig,
|
||||
)
|
||||
from ray.rllib.connectors.common.add_observations_from_episodes_to_batch import (
|
||||
AddObservationsFromEpisodesToBatch,
|
||||
)
|
||||
from ray.rllib.connectors.learner.add_next_observations_from_episodes_to_train_batch import ( # noqa
|
||||
AddNextObservationsFromEpisodesToTrainBatch,
|
||||
)
|
||||
from ray.rllib.core.learner.learner import Learner
|
||||
from ray.rllib.core.rl_module.rl_module import RLModuleSpec
|
||||
from ray.rllib.execution.rollout_ops import (
|
||||
synchronous_parallel_sample,
|
||||
)
|
||||
from ray.rllib.execution.train_ops import (
|
||||
multi_gpu_train_one_step,
|
||||
train_one_step,
|
||||
)
|
||||
from ray.rllib.policy.policy import Policy
|
||||
from ray.rllib.utils.annotations import OldAPIStack, override
|
||||
from ray.rllib.utils.framework import try_import_tf, try_import_tfp
|
||||
from ray.rllib.utils.metrics import (
|
||||
LAST_TARGET_UPDATE_TS,
|
||||
LEARNER_RESULTS,
|
||||
LEARNER_UPDATE_TIMER,
|
||||
NUM_AGENT_STEPS_SAMPLED,
|
||||
NUM_AGENT_STEPS_TRAINED,
|
||||
NUM_ENV_STEPS_SAMPLED,
|
||||
NUM_ENV_STEPS_TRAINED,
|
||||
NUM_TARGET_UPDATES,
|
||||
OFFLINE_SAMPLING_TIMER,
|
||||
SAMPLE_TIMER,
|
||||
SYNCH_WORKER_WEIGHTS_TIMER,
|
||||
TARGET_NET_UPDATE_TIMER,
|
||||
TIMERS,
|
||||
)
|
||||
from ray.rllib.utils.typing import ResultDict, RLModuleSpecType
|
||||
|
||||
tf1, tf, tfv = try_import_tf()
|
||||
tfp = try_import_tfp()
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class CQLConfig(SACConfig):
|
||||
"""Defines a configuration class from which a CQL can be built.
|
||||
|
||||
.. testcode::
|
||||
:skipif: True
|
||||
|
||||
from ray.rllib.algorithms.cql import CQLConfig
|
||||
config = CQLConfig().training(gamma=0.9, lr=0.01)
|
||||
config = config.resources(num_gpus=0)
|
||||
config = config.env_runners(num_env_runners=4)
|
||||
print(config.to_dict())
|
||||
# Build a Algorithm object from the config and run 1 training iteration.
|
||||
algo = config.build(env="CartPole-v1")
|
||||
algo.train()
|
||||
"""
|
||||
|
||||
def __init__(self, algo_class=None):
|
||||
super().__init__(algo_class=algo_class or CQL)
|
||||
|
||||
# fmt: off
|
||||
# __sphinx_doc_begin__
|
||||
# CQL-specific config settings:
|
||||
self.bc_iters = 20000
|
||||
self.temperature = 1.0
|
||||
self.num_actions = 10
|
||||
self.lagrangian = False
|
||||
self.lagrangian_thresh = 5.0
|
||||
self.min_q_weight = 5.0
|
||||
self.deterministic_backup = True
|
||||
self.lr = 3e-4
|
||||
# Note, the new stack defines learning rates for each component.
|
||||
# The base learning rate `lr` has to be set to `None`, if using
|
||||
# the new stack.
|
||||
self.actor_lr = 1e-4
|
||||
self.critic_lr = 1e-3
|
||||
self.alpha_lr = 1e-3
|
||||
|
||||
self.replay_buffer_config = {
|
||||
"_enable_replay_buffer_api": True,
|
||||
"type": "MultiAgentPrioritizedReplayBuffer",
|
||||
"capacity": int(1e6),
|
||||
# If True prioritized replay buffer will be used.
|
||||
"prioritized_replay": False,
|
||||
"prioritized_replay_alpha": 0.6,
|
||||
"prioritized_replay_beta": 0.4,
|
||||
"prioritized_replay_eps": 1e-6,
|
||||
# Whether to compute priorities already on the remote worker side.
|
||||
"worker_side_prioritization": False,
|
||||
}
|
||||
|
||||
# Changes to Algorithm's/SACConfig's default:
|
||||
|
||||
# .reporting()
|
||||
self.min_sample_timesteps_per_iteration = 0
|
||||
self.min_train_timesteps_per_iteration = 100
|
||||
# fmt: on
|
||||
# __sphinx_doc_end__
|
||||
|
||||
self.timesteps_per_iteration = DEPRECATED_VALUE
|
||||
|
||||
@override(SACConfig)
|
||||
def training(
|
||||
self,
|
||||
*,
|
||||
bc_iters: Optional[int] = NotProvided,
|
||||
temperature: Optional[float] = NotProvided,
|
||||
num_actions: Optional[int] = NotProvided,
|
||||
lagrangian: Optional[bool] = NotProvided,
|
||||
lagrangian_thresh: Optional[float] = NotProvided,
|
||||
min_q_weight: Optional[float] = NotProvided,
|
||||
deterministic_backup: Optional[bool] = NotProvided,
|
||||
**kwargs,
|
||||
) -> Self:
|
||||
"""Sets the training-related configuration.
|
||||
|
||||
Args:
|
||||
bc_iters: Number of iterations with Behavior Cloning pretraining.
|
||||
temperature: CQL loss temperature.
|
||||
num_actions: Number of actions to sample for CQL loss
|
||||
lagrangian: Whether to use the Lagrangian for Alpha Prime (in CQL loss).
|
||||
lagrangian_thresh: Lagrangian threshold.
|
||||
min_q_weight: in Q weight multiplier.
|
||||
deterministic_backup: If the target in the Bellman update should have an
|
||||
entropy backup. Defaults to `True`.
|
||||
|
||||
Returns:
|
||||
This updated AlgorithmConfig object.
|
||||
"""
|
||||
# Pass kwargs onto super's `training()` method.
|
||||
super().training(**kwargs)
|
||||
|
||||
if bc_iters is not NotProvided:
|
||||
self.bc_iters = bc_iters
|
||||
if temperature is not NotProvided:
|
||||
self.temperature = temperature
|
||||
if num_actions is not NotProvided:
|
||||
self.num_actions = num_actions
|
||||
if lagrangian is not NotProvided:
|
||||
self.lagrangian = lagrangian
|
||||
if lagrangian_thresh is not NotProvided:
|
||||
self.lagrangian_thresh = lagrangian_thresh
|
||||
if min_q_weight is not NotProvided:
|
||||
self.min_q_weight = min_q_weight
|
||||
if deterministic_backup is not NotProvided:
|
||||
self.deterministic_backup = deterministic_backup
|
||||
|
||||
return self
|
||||
|
||||
@override(AlgorithmConfig)
|
||||
def offline_data(self, **kwargs) -> Self:
|
||||
|
||||
super().offline_data(**kwargs)
|
||||
|
||||
# Check, if the passed in class incorporates the `OfflinePreLearner`
|
||||
# interface.
|
||||
if "prelearner_class" in kwargs:
|
||||
from ray.rllib.offline.offline_data import OfflinePreLearner
|
||||
|
||||
if not issubclass(kwargs.get("prelearner_class"), OfflinePreLearner):
|
||||
raise ValueError(
|
||||
f"`prelearner_class` {kwargs.get('prelearner_class')} is not a "
|
||||
"subclass of `OfflinePreLearner`. Any class passed to "
|
||||
"`prelearner_class` needs to implement the interface given by "
|
||||
"`OfflinePreLearner`."
|
||||
)
|
||||
|
||||
return self
|
||||
|
||||
@override(SACConfig)
|
||||
def get_default_learner_class(self) -> Union[Type["Learner"], str]:
|
||||
if self.framework_str == "torch":
|
||||
from ray.rllib.algorithms.cql.torch.cql_torch_learner import CQLTorchLearner
|
||||
|
||||
return CQLTorchLearner
|
||||
else:
|
||||
raise ValueError(
|
||||
f"The framework {self.framework_str} is not supported. "
|
||||
"Use `'torch'` instead."
|
||||
)
|
||||
|
||||
@override(AlgorithmConfig)
|
||||
def build_learner_connector(
|
||||
self,
|
||||
input_observation_space,
|
||||
input_action_space,
|
||||
device=None,
|
||||
):
|
||||
pipeline = super().build_learner_connector(
|
||||
input_observation_space=input_observation_space,
|
||||
input_action_space=input_action_space,
|
||||
device=device,
|
||||
)
|
||||
|
||||
# Prepend the "add-NEXT_OBS-from-episodes-to-train-batch" connector piece (right
|
||||
# after the corresponding "add-OBS-..." default piece).
|
||||
pipeline.insert_after(
|
||||
AddObservationsFromEpisodesToBatch,
|
||||
AddNextObservationsFromEpisodesToTrainBatch(),
|
||||
)
|
||||
|
||||
return pipeline
|
||||
|
||||
@override(SACConfig)
|
||||
def validate(self) -> None:
|
||||
# First check, whether old `timesteps_per_iteration` is used.
|
||||
if self.timesteps_per_iteration != DEPRECATED_VALUE:
|
||||
deprecation_warning(
|
||||
old="timesteps_per_iteration",
|
||||
new="min_train_timesteps_per_iteration",
|
||||
error=True,
|
||||
)
|
||||
|
||||
# Call super's validation method.
|
||||
super().validate()
|
||||
|
||||
# CQL-torch performs the optimizer steps inside the loss function.
|
||||
# Using the multi-GPU optimizer will therefore not work (see multi-GPU
|
||||
# check above) and we must use the simple optimizer for now.
|
||||
if self.simple_optimizer is not True and self.framework_str == "torch":
|
||||
self.simple_optimizer = True
|
||||
|
||||
if self.framework_str in ["tf", "tf2"] and tfp is None:
|
||||
logger.warning(
|
||||
"You need `tensorflow_probability` in order to run CQL! "
|
||||
"Install it via `pip install tensorflow_probability`. Your "
|
||||
f"tf.__version__={tf.__version__ if tf else None}."
|
||||
"Trying to import tfp results in the following error:"
|
||||
)
|
||||
try_import_tfp(error=True)
|
||||
|
||||
# Assert that for a local learner the number of iterations is 1. Note,
|
||||
# this is needed because we have no iterators, but instead a single
|
||||
# batch returned directly from the `OfflineData.sample` method.
|
||||
if (
|
||||
self.num_learners == 0
|
||||
and not self.dataset_num_iters_per_learner
|
||||
and self.enable_rl_module_and_learner
|
||||
):
|
||||
self._value_error(
|
||||
"When using a single local learner the number of iterations "
|
||||
"per learner, `dataset_num_iters_per_learner` has to be defined. "
|
||||
"Set this hyperparameter in the `AlgorithmConfig.offline_data`."
|
||||
)
|
||||
|
||||
@override(SACConfig)
|
||||
def get_default_rl_module_spec(self) -> RLModuleSpecType:
|
||||
if self.framework_str == "torch":
|
||||
from ray.rllib.algorithms.cql.torch.default_cql_torch_rl_module import (
|
||||
DefaultCQLTorchRLModule,
|
||||
)
|
||||
|
||||
return RLModuleSpec(module_class=DefaultCQLTorchRLModule)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"The framework {self.framework_str} is not supported. Use `torch`."
|
||||
)
|
||||
|
||||
@property
|
||||
def _model_config_auto_includes(self):
|
||||
return super()._model_config_auto_includes | {
|
||||
"num_actions": self.num_actions,
|
||||
}
|
||||
|
||||
|
||||
class CQL(SAC):
|
||||
"""CQL (derived from SAC)."""
|
||||
|
||||
@classmethod
|
||||
@override(SAC)
|
||||
def get_default_config(cls) -> CQLConfig:
|
||||
return CQLConfig()
|
||||
|
||||
@classmethod
|
||||
@override(SAC)
|
||||
def get_default_policy_class(
|
||||
cls, config: AlgorithmConfig
|
||||
) -> Optional[Type[Policy]]:
|
||||
if config["framework"] == "torch":
|
||||
return CQLTorchPolicy
|
||||
else:
|
||||
return CQLTFPolicy
|
||||
|
||||
@override(SAC)
|
||||
def training_step(self) -> None:
|
||||
# Old API stack (Policy, RolloutWorker, Connector).
|
||||
if not self.config.enable_env_runner_and_connector_v2:
|
||||
return self._training_step_old_api_stack()
|
||||
|
||||
# Sampling from offline data.
|
||||
with self.metrics.log_time((TIMERS, OFFLINE_SAMPLING_TIMER)):
|
||||
# If we should use an iterator in the learner(s). Note, in case of
|
||||
# multiple learners we must always return a list of iterators.
|
||||
return_iterator = return_iterator = (
|
||||
self.config.num_learners > 0
|
||||
or self.config.dataset_num_iters_per_learner != 1
|
||||
)
|
||||
|
||||
# Return an iterator in case we are using remote learners.
|
||||
batch_or_iterator = self.offline_data.sample(
|
||||
num_samples=self.config.train_batch_size_per_learner,
|
||||
num_shards=self.config.num_learners,
|
||||
# Return an iterator, if a `Learner` should update
|
||||
# multiple times per RLlib iteration.
|
||||
return_iterator=return_iterator,
|
||||
)
|
||||
|
||||
# Updating the policy.
|
||||
with self.metrics.log_time((TIMERS, LEARNER_UPDATE_TIMER)):
|
||||
learner_results = self.learner_group.update(
|
||||
data_iterators=batch_or_iterator,
|
||||
minibatch_size=self.config.train_batch_size_per_learner,
|
||||
num_iters=self.config.dataset_num_iters_per_learner,
|
||||
)
|
||||
|
||||
# Log training results.
|
||||
self.metrics.aggregate(learner_results, key=LEARNER_RESULTS)
|
||||
|
||||
@OldAPIStack
|
||||
def _training_step_old_api_stack(self) -> ResultDict:
|
||||
# Collect SampleBatches from sample workers.
|
||||
with self._timers[SAMPLE_TIMER]:
|
||||
train_batch = synchronous_parallel_sample(worker_set=self.env_runner_group)
|
||||
train_batch = train_batch.as_multi_agent()
|
||||
self._counters[NUM_AGENT_STEPS_SAMPLED] += train_batch.agent_steps()
|
||||
self._counters[NUM_ENV_STEPS_SAMPLED] += train_batch.env_steps()
|
||||
|
||||
# Postprocess batch before we learn on it.
|
||||
post_fn = self.config.get("before_learn_on_batch") or (lambda b, *a: b)
|
||||
train_batch = post_fn(train_batch, self.env_runner_group, self.config)
|
||||
|
||||
# Learn on training batch.
|
||||
# Use simple optimizer (only for multi-agent or tf-eager; all other
|
||||
# cases should use the multi-GPU optimizer, even if only using 1 GPU)
|
||||
if self.config.get("simple_optimizer") is True:
|
||||
train_results = train_one_step(self, train_batch)
|
||||
else:
|
||||
train_results = multi_gpu_train_one_step(self, train_batch)
|
||||
|
||||
# Update target network every `target_network_update_freq` training steps.
|
||||
cur_ts = self._counters[
|
||||
NUM_AGENT_STEPS_TRAINED
|
||||
if self.config.count_steps_by == "agent_steps"
|
||||
else NUM_ENV_STEPS_TRAINED
|
||||
]
|
||||
last_update = self._counters[LAST_TARGET_UPDATE_TS]
|
||||
if cur_ts - last_update >= self.config.target_network_update_freq:
|
||||
with self._timers[TARGET_NET_UPDATE_TIMER]:
|
||||
to_update = self.env_runner.get_policies_to_train()
|
||||
self.env_runner.foreach_policy_to_train(
|
||||
lambda p, pid: pid in to_update and p.update_target()
|
||||
)
|
||||
self._counters[NUM_TARGET_UPDATES] += 1
|
||||
self._counters[LAST_TARGET_UPDATE_TS] = cur_ts
|
||||
|
||||
# Update remote workers's weights after learning on local worker
|
||||
# (only those policies that were actually trained).
|
||||
if self.env_runner_group.num_remote_workers() > 0:
|
||||
with self._timers[SYNCH_WORKER_WEIGHTS_TIMER]:
|
||||
self.env_runner_group.sync_weights(policies=list(train_results.keys()))
|
||||
|
||||
# Return all collected metrics for the iteration.
|
||||
return train_results
|
||||
@@ -0,0 +1,426 @@
|
||||
"""
|
||||
TensorFlow policy class used for CQL.
|
||||
"""
|
||||
import logging
|
||||
from functools import partial
|
||||
from typing import Dict, List, Type, Union
|
||||
|
||||
import gymnasium as gym
|
||||
import numpy as np
|
||||
import tree
|
||||
|
||||
import ray
|
||||
from ray.rllib.algorithms.sac.sac_tf_policy import (
|
||||
ActorCriticOptimizerMixin as SACActorCriticOptimizerMixin,
|
||||
ComputeTDErrorMixin,
|
||||
_get_dist_class,
|
||||
apply_gradients as sac_apply_gradients,
|
||||
build_sac_model,
|
||||
compute_and_clip_gradients as sac_compute_and_clip_gradients,
|
||||
get_distribution_inputs_and_class,
|
||||
postprocess_trajectory,
|
||||
setup_late_mixins,
|
||||
stats,
|
||||
validate_spaces,
|
||||
)
|
||||
from ray.rllib.models.modelv2 import ModelV2
|
||||
from ray.rllib.models.tf.tf_action_dist import TFActionDistribution
|
||||
from ray.rllib.policy.policy import Policy
|
||||
from ray.rllib.policy.sample_batch import SampleBatch
|
||||
from ray.rllib.policy.tf_mixins import TargetNetworkMixin
|
||||
from ray.rllib.policy.tf_policy_template import build_tf_policy
|
||||
from ray.rllib.utils.exploration.random import Random
|
||||
from ray.rllib.utils.framework import get_variable, try_import_tf, try_import_tfp
|
||||
from ray.rllib.utils.typing import (
|
||||
AlgorithmConfigDict,
|
||||
LocalOptimizer,
|
||||
ModelGradients,
|
||||
TensorType,
|
||||
)
|
||||
|
||||
tf1, tf, tfv = try_import_tf()
|
||||
tfp = try_import_tfp()
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
MEAN_MIN = -9.0
|
||||
MEAN_MAX = 9.0
|
||||
|
||||
|
||||
def _repeat_tensor(t: TensorType, n: int):
|
||||
# Insert new axis at position 1 into tensor t
|
||||
t_rep = tf.expand_dims(t, 1)
|
||||
# Repeat tensor t_rep along new axis n times
|
||||
multiples = tf.concat([[1, n], tf.tile([1], tf.expand_dims(tf.rank(t) - 1, 0))], 0)
|
||||
t_rep = tf.tile(t_rep, multiples)
|
||||
# Merge new axis into batch axis
|
||||
t_rep = tf.reshape(t_rep, tf.concat([[-1], tf.shape(t)[1:]], 0))
|
||||
return t_rep
|
||||
|
||||
|
||||
# Returns policy tiled actions and log probabilities for CQL Loss
|
||||
def policy_actions_repeat(model, action_dist, obs, num_repeat=1):
|
||||
batch_size = tf.shape(tree.flatten(obs)[0])[0]
|
||||
obs_temp = tree.map_structure(lambda t: _repeat_tensor(t, num_repeat), obs)
|
||||
logits, _ = model.get_action_model_outputs(obs_temp)
|
||||
policy_dist = action_dist(logits, model)
|
||||
actions, logp_ = policy_dist.sample_logp()
|
||||
logp = tf.expand_dims(logp_, -1)
|
||||
return actions, tf.reshape(logp, [batch_size, num_repeat, 1])
|
||||
|
||||
|
||||
def q_values_repeat(model, obs, actions, twin=False):
|
||||
action_shape = tf.shape(actions)[0]
|
||||
obs_shape = tf.shape(tree.flatten(obs)[0])[0]
|
||||
num_repeat = action_shape // obs_shape
|
||||
obs_temp = tree.map_structure(lambda t: _repeat_tensor(t, num_repeat), obs)
|
||||
if not twin:
|
||||
preds_, _ = model.get_q_values(obs_temp, actions)
|
||||
else:
|
||||
preds_, _ = model.get_twin_q_values(obs_temp, actions)
|
||||
preds = tf.reshape(preds_, [obs_shape, num_repeat, 1])
|
||||
return preds
|
||||
|
||||
|
||||
def cql_loss(
|
||||
policy: Policy,
|
||||
model: ModelV2,
|
||||
dist_class: Type[TFActionDistribution],
|
||||
train_batch: SampleBatch,
|
||||
) -> Union[TensorType, List[TensorType]]:
|
||||
logger.info(f"Current iteration = {policy.cur_iter}")
|
||||
policy.cur_iter += 1
|
||||
|
||||
# For best performance, turn deterministic off
|
||||
deterministic = policy.config["_deterministic_loss"]
|
||||
assert not deterministic
|
||||
twin_q = policy.config["twin_q"]
|
||||
discount = policy.config["gamma"]
|
||||
|
||||
# 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 = tf.cast(train_batch[SampleBatch.ACTIONS], tf.float32)
|
||||
rewards = tf.cast(train_batch[SampleBatch.REWARDS], tf.float32)
|
||||
next_obs = train_batch[SampleBatch.NEXT_OBS]
|
||||
terminals = train_batch[SampleBatch.TERMINATEDS]
|
||||
|
||||
model_out_t, _ = model(SampleBatch(obs=obs, _is_training=True), [], None)
|
||||
|
||||
model_out_tp1, _ = model(SampleBatch(obs=next_obs, _is_training=True), [], None)
|
||||
|
||||
target_model_out_tp1, _ = policy.target_model(
|
||||
SampleBatch(obs=next_obs, _is_training=True), [], None
|
||||
)
|
||||
|
||||
action_dist_class = _get_dist_class(policy, policy.config, policy.action_space)
|
||||
action_dist_inputs_t, _ = model.get_action_model_outputs(model_out_t)
|
||||
action_dist_t = action_dist_class(action_dist_inputs_t, model)
|
||||
policy_t, log_pis_t = action_dist_t.sample_logp()
|
||||
log_pis_t = tf.expand_dims(log_pis_t, -1)
|
||||
|
||||
# Unlike original SAC, Alpha and Actor Loss are computed first.
|
||||
# Alpha Loss
|
||||
alpha_loss = -tf.reduce_mean(
|
||||
model.log_alpha * tf.stop_gradient(log_pis_t + model.target_entropy)
|
||||
)
|
||||
|
||||
# Policy Loss (Either Behavior Clone Loss or SAC Loss)
|
||||
alpha = tf.math.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 = tf.math.minimum(min_q, twin_q_)
|
||||
actor_loss = tf.reduce_mean(tf.stop_gradient(alpha) * log_pis_t - min_q)
|
||||
else:
|
||||
bc_logp = action_dist_t.logp(actions)
|
||||
actor_loss = tf.reduce_mean(tf.stop_gradient(alpha) * log_pis_t - bc_logp)
|
||||
# actor_loss = -tf.reduce_mean(bc_logp)
|
||||
|
||||
# Critic Loss (Standard SAC Critic L2 Loss + CQL Entropy Loss)
|
||||
# SAC Loss:
|
||||
# Q-values for the batched actions.
|
||||
action_dist_inputs_tp1, _ = model.get_action_model_outputs(model_out_tp1)
|
||||
action_dist_tp1 = action_dist_class(action_dist_inputs_tp1, model)
|
||||
policy_tp1, _ = action_dist_tp1.sample_logp()
|
||||
|
||||
q_t, _ = model.get_q_values(model_out_t, actions)
|
||||
q_t_selected = tf.squeeze(q_t, axis=-1)
|
||||
if twin_q:
|
||||
twin_q_t, _ = model.get_twin_q_values(model_out_t, actions)
|
||||
twin_q_t_selected = tf.squeeze(twin_q_t, axis=-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 = tf.math.minimum(q_tp1, twin_q_tp1)
|
||||
|
||||
q_tp1_best = tf.squeeze(input=q_tp1, axis=-1)
|
||||
q_tp1_best_masked = (1.0 - tf.cast(terminals, tf.float32)) * q_tp1_best
|
||||
|
||||
# compute RHS of bellman equation
|
||||
q_t_target = tf.stop_gradient(
|
||||
rewards + (discount ** policy.config["n_step"]) * q_tp1_best_masked
|
||||
)
|
||||
|
||||
# Compute the TD-error (potentially clipped), for priority replay buffer
|
||||
base_td_error = tf.math.abs(q_t_selected - q_t_target)
|
||||
if twin_q:
|
||||
twin_td_error = tf.math.abs(twin_q_t_selected - q_t_target)
|
||||
td_error = 0.5 * (base_td_error + twin_td_error)
|
||||
else:
|
||||
td_error = base_td_error
|
||||
|
||||
critic_loss_1 = tf.keras.losses.MSE(q_t_selected, q_t_target)
|
||||
if twin_q:
|
||||
critic_loss_2 = tf.keras.losses.MSE(twin_q_t_selected, q_t_target)
|
||||
|
||||
# CQL Loss (We are using Entropy version of CQL (the best version))
|
||||
rand_actions, _ = policy._random_action_generator.get_exploration_action(
|
||||
action_distribution=action_dist_class(
|
||||
tf.tile(action_dist_tp1.inputs, (num_actions, 1)), model
|
||||
),
|
||||
timestep=0,
|
||||
explore=True,
|
||||
)
|
||||
curr_actions, curr_logp = policy_actions_repeat(
|
||||
model, action_dist_class, model_out_t, num_actions
|
||||
)
|
||||
next_actions, next_logp = policy_actions_repeat(
|
||||
model, action_dist_class, model_out_tp1, num_actions
|
||||
)
|
||||
|
||||
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 ** int(curr_actions.shape[-1]))
|
||||
cat_q1 = tf.concat(
|
||||
[
|
||||
q1_rand - random_density,
|
||||
q1_next_actions - tf.stop_gradient(next_logp),
|
||||
q1_curr_actions - tf.stop_gradient(curr_logp),
|
||||
],
|
||||
1,
|
||||
)
|
||||
if twin_q:
|
||||
cat_q2 = tf.concat(
|
||||
[
|
||||
q2_rand - random_density,
|
||||
q2_next_actions - tf.stop_gradient(next_logp),
|
||||
q2_curr_actions - tf.stop_gradient(curr_logp),
|
||||
],
|
||||
1,
|
||||
)
|
||||
|
||||
min_qf1_loss_ = (
|
||||
tf.reduce_mean(tf.reduce_logsumexp(cat_q1 / cql_temp, axis=1))
|
||||
* min_q_weight
|
||||
* cql_temp
|
||||
)
|
||||
min_qf1_loss = min_qf1_loss_ - (tf.reduce_mean(q_t) * min_q_weight)
|
||||
if twin_q:
|
||||
min_qf2_loss_ = (
|
||||
tf.reduce_mean(tf.reduce_logsumexp(cat_q2 / cql_temp, axis=1))
|
||||
* min_q_weight
|
||||
* cql_temp
|
||||
)
|
||||
min_qf2_loss = min_qf2_loss_ - (tf.reduce_mean(twin_q_t) * min_q_weight)
|
||||
|
||||
if use_lagrange:
|
||||
alpha_prime = tf.clip_by_value(model.log_alpha_prime.exp(), 0.0, 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_qf1_loss]
|
||||
if twin_q:
|
||||
cql_loss.append(min_qf2_loss)
|
||||
|
||||
critic_loss = [critic_loss_1 + min_qf1_loss]
|
||||
if twin_q:
|
||||
critic_loss.append(critic_loss_2 + min_qf2_loss)
|
||||
|
||||
# Save for stats function.
|
||||
policy.q_t = q_t_selected
|
||||
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 actor_loss + tf.math.add_n(critic_loss) + alpha_loss + alpha_prime_loss
|
||||
return actor_loss + tf.math.add_n(critic_loss) + alpha_loss
|
||||
|
||||
|
||||
def cql_stats(policy: Policy, train_batch: SampleBatch) -> Dict[str, TensorType]:
|
||||
sac_dict = stats(policy, train_batch)
|
||||
sac_dict["cql_loss"] = tf.reduce_mean(tf.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
|
||||
|
||||
|
||||
class ActorCriticOptimizerMixin(SACActorCriticOptimizerMixin):
|
||||
def __init__(self, config):
|
||||
super().__init__(config)
|
||||
if config["lagrangian"]:
|
||||
# Eager mode.
|
||||
if config["framework"] == "tf2":
|
||||
self._alpha_prime_optimizer = tf.keras.optimizers.Adam(
|
||||
learning_rate=config["optimization"]["critic_learning_rate"]
|
||||
)
|
||||
# Static graph mode.
|
||||
else:
|
||||
self._alpha_prime_optimizer = tf1.train.AdamOptimizer(
|
||||
learning_rate=config["optimization"]["critic_learning_rate"]
|
||||
)
|
||||
|
||||
|
||||
def setup_early_mixins(
|
||||
policy: Policy,
|
||||
obs_space: gym.spaces.Space,
|
||||
action_space: gym.spaces.Space,
|
||||
config: AlgorithmConfigDict,
|
||||
) -> None:
|
||||
"""Call mixin classes' constructors before Policy's initialization.
|
||||
|
||||
Adds the necessary optimizers to the given Policy.
|
||||
|
||||
Args:
|
||||
policy: The Policy object.
|
||||
obs_space (gym.spaces.Space): The Policy's observation space.
|
||||
action_space (gym.spaces.Space): The Policy's action space.
|
||||
config: The Policy's config.
|
||||
"""
|
||||
policy.cur_iter = 0
|
||||
ActorCriticOptimizerMixin.__init__(policy, config)
|
||||
if config["lagrangian"]:
|
||||
policy.model.log_alpha_prime = get_variable(
|
||||
0.0, framework="tf", trainable=True, tf_name="log_alpha_prime"
|
||||
)
|
||||
policy.alpha_prime_optim = tf.keras.optimizers.Adam(
|
||||
learning_rate=config["optimization"]["critic_learning_rate"],
|
||||
)
|
||||
# Generic random action generator for calculating CQL-loss.
|
||||
policy._random_action_generator = Random(
|
||||
action_space,
|
||||
model=None,
|
||||
framework="tf2",
|
||||
policy_config=config,
|
||||
num_workers=0,
|
||||
worker_index=0,
|
||||
)
|
||||
|
||||
|
||||
def compute_gradients_fn(
|
||||
policy: Policy, optimizer: LocalOptimizer, loss: TensorType
|
||||
) -> ModelGradients:
|
||||
grads_and_vars = sac_compute_and_clip_gradients(policy, optimizer, loss)
|
||||
|
||||
if policy.config["lagrangian"]:
|
||||
# Eager: Use GradientTape (which is a property of the `optimizer`
|
||||
# object (an OptimizerWrapper): see rllib/policy/eager_tf_policy.py).
|
||||
if policy.config["framework"] == "tf2":
|
||||
tape = optimizer.tape
|
||||
log_alpha_prime = [policy.model.log_alpha_prime]
|
||||
alpha_prime_grads_and_vars = list(
|
||||
zip(
|
||||
tape.gradient(policy.alpha_prime_loss, log_alpha_prime),
|
||||
log_alpha_prime,
|
||||
)
|
||||
)
|
||||
# Tf1.x: Use optimizer.compute_gradients()
|
||||
else:
|
||||
alpha_prime_grads_and_vars = (
|
||||
policy._alpha_prime_optimizer.compute_gradients(
|
||||
policy.alpha_prime_loss, var_list=[policy.model.log_alpha_prime]
|
||||
)
|
||||
)
|
||||
|
||||
# Clip if necessary.
|
||||
if policy.config["grad_clip"]:
|
||||
clip_func = partial(tf.clip_by_norm, clip_norm=policy.config["grad_clip"])
|
||||
else:
|
||||
clip_func = tf.identity
|
||||
|
||||
# Save grads and vars for later use in `build_apply_op`.
|
||||
policy._alpha_prime_grads_and_vars = [
|
||||
(clip_func(g), v) for (g, v) in alpha_prime_grads_and_vars if g is not None
|
||||
]
|
||||
|
||||
grads_and_vars += policy._alpha_prime_grads_and_vars
|
||||
return grads_and_vars
|
||||
|
||||
|
||||
def apply_gradients_fn(policy, optimizer, grads_and_vars):
|
||||
sac_results = sac_apply_gradients(policy, optimizer, grads_and_vars)
|
||||
|
||||
if policy.config["lagrangian"]:
|
||||
# Eager mode -> Just apply and return None.
|
||||
if policy.config["framework"] == "tf2":
|
||||
policy._alpha_prime_optimizer.apply_gradients(
|
||||
policy._alpha_prime_grads_and_vars
|
||||
)
|
||||
return
|
||||
# Tf static graph -> Return grouped op.
|
||||
else:
|
||||
alpha_prime_apply_op = policy._alpha_prime_optimizer.apply_gradients(
|
||||
policy._alpha_prime_grads_and_vars,
|
||||
global_step=tf1.train.get_or_create_global_step(),
|
||||
)
|
||||
return tf.group([sac_results, alpha_prime_apply_op])
|
||||
return sac_results
|
||||
|
||||
|
||||
# Build a child class of `TFPolicy`, given the custom functions defined
|
||||
# above.
|
||||
CQLTFPolicy = build_tf_policy(
|
||||
name="CQLTFPolicy",
|
||||
loss_fn=cql_loss,
|
||||
get_default_config=lambda: ray.rllib.algorithms.cql.cql.CQLConfig(),
|
||||
validate_spaces=validate_spaces,
|
||||
stats_fn=cql_stats,
|
||||
postprocess_fn=postprocess_trajectory,
|
||||
before_init=setup_early_mixins,
|
||||
after_init=setup_late_mixins,
|
||||
make_model=build_sac_model,
|
||||
extra_learn_fetches_fn=lambda policy: {"td_error": policy.td_error},
|
||||
mixins=[ActorCriticOptimizerMixin, TargetNetworkMixin, ComputeTDErrorMixin],
|
||||
action_distribution_fn=get_distribution_inputs_and_class,
|
||||
compute_gradients_fn=compute_gradients_fn,
|
||||
apply_gradients_fn=apply_gradients_fn,
|
||||
)
|
||||
@@ -0,0 +1,406 @@
|
||||
"""
|
||||
PyTorch policy class used for CQL.
|
||||
"""
|
||||
import logging
|
||||
from typing import Dict, List, Tuple, Type, Union
|
||||
|
||||
import gymnasium as gym
|
||||
import numpy as np
|
||||
import tree
|
||||
|
||||
import ray
|
||||
from ray.rllib.algorithms.sac.sac_tf_policy import (
|
||||
postprocess_trajectory,
|
||||
validate_spaces,
|
||||
)
|
||||
from ray.rllib.algorithms.sac.sac_torch_policy import (
|
||||
ComputeTDErrorMixin,
|
||||
_get_dist_class,
|
||||
action_distribution_fn,
|
||||
build_sac_model_and_action_dist,
|
||||
optimizer_fn,
|
||||
setup_late_mixins,
|
||||
stats,
|
||||
)
|
||||
from ray.rllib.models.modelv2 import ModelV2
|
||||
from ray.rllib.models.torch.torch_action_dist import TorchDistributionWrapper
|
||||
from ray.rllib.policy.policy import Policy
|
||||
from ray.rllib.policy.policy_template import build_policy_class
|
||||
from ray.rllib.policy.sample_batch import SampleBatch
|
||||
from ray.rllib.policy.torch_mixins import TargetNetworkMixin
|
||||
from ray.rllib.utils.framework import try_import_torch
|
||||
from ray.rllib.utils.metrics.learner_info import LEARNER_STATS_KEY
|
||||
from ray.rllib.utils.torch_utils import (
|
||||
apply_grad_clipping,
|
||||
concat_multi_gpu_td_errors,
|
||||
convert_to_torch_tensor,
|
||||
)
|
||||
from ray.rllib.utils.typing import AlgorithmConfigDict, LocalOptimizer, TensorType
|
||||
|
||||
torch, nn = try_import_torch()
|
||||
F = nn.functional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
MEAN_MIN = -9.0
|
||||
MEAN_MAX = 9.0
|
||||
|
||||
|
||||
def _repeat_tensor(t: TensorType, n: int):
|
||||
# Insert new dimension at posotion 1 into tensor t
|
||||
t_rep = t.unsqueeze(1)
|
||||
# Repeat tensor t_rep along new dimension n times
|
||||
t_rep = torch.repeat_interleave(t_rep, n, dim=1)
|
||||
# Merge new dimension into batch dimension
|
||||
t_rep = t_rep.view(-1, *t.shape[1:])
|
||||
return t_rep
|
||||
|
||||
|
||||
# Returns policy tiled actions and log probabilities for CQL Loss
|
||||
def policy_actions_repeat(model, action_dist, obs, num_repeat=1):
|
||||
batch_size = tree.flatten(obs)[0].shape[0]
|
||||
obs_temp = tree.map_structure(lambda t: _repeat_tensor(t, num_repeat), obs)
|
||||
logits, _ = model.get_action_model_outputs(obs_temp)
|
||||
policy_dist = action_dist(logits, model)
|
||||
actions, logp_ = policy_dist.sample_logp()
|
||||
logp = logp_.unsqueeze(-1)
|
||||
return actions, logp.view(batch_size, num_repeat, 1)
|
||||
|
||||
|
||||
def q_values_repeat(model, obs, actions, twin=False):
|
||||
action_shape = actions.shape[0]
|
||||
obs_shape = tree.flatten(obs)[0].shape[0]
|
||||
num_repeat = int(action_shape / obs_shape)
|
||||
obs_temp = tree.map_structure(lambda t: _repeat_tensor(t, num_repeat), obs)
|
||||
if not twin:
|
||||
preds_, _ = model.get_q_values(obs_temp, actions)
|
||||
else:
|
||||
preds_, _ = model.get_twin_q_values(obs_temp, actions)
|
||||
preds = preds_.view(obs_shape, num_repeat, 1)
|
||||
return preds
|
||||
|
||||
|
||||
def cql_loss(
|
||||
policy: Policy,
|
||||
model: ModelV2,
|
||||
dist_class: Type[TorchDistributionWrapper],
|
||||
train_batch: SampleBatch,
|
||||
) -> Union[TensorType, List[TensorType]]:
|
||||
logger.info(f"Current iteration = {policy.cur_iter}")
|
||||
policy.cur_iter += 1
|
||||
|
||||
# Look up the target model (tower) using the model tower.
|
||||
target_model = policy.target_models[model]
|
||||
|
||||
# For best performance, turn deterministic off
|
||||
deterministic = policy.config["_deterministic_loss"]
|
||||
assert not deterministic
|
||||
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].float()
|
||||
next_obs = train_batch[SampleBatch.NEXT_OBS]
|
||||
terminals = train_batch[SampleBatch.TERMINATEDS]
|
||||
|
||||
model_out_t, _ = model(SampleBatch(obs=obs, _is_training=True), [], None)
|
||||
|
||||
model_out_tp1, _ = model(SampleBatch(obs=next_obs, _is_training=True), [], None)
|
||||
|
||||
target_model_out_tp1, _ = target_model(
|
||||
SampleBatch(obs=next_obs, _is_training=True), [], None
|
||||
)
|
||||
|
||||
action_dist_class = _get_dist_class(policy, policy.config, policy.action_space)
|
||||
action_dist_inputs_t, _ = model.get_action_model_outputs(model_out_t)
|
||||
action_dist_t = action_dist_class(action_dist_inputs_t, model)
|
||||
policy_t, log_pis_t = action_dist_t.sample_logp()
|
||||
log_pis_t = torch.unsqueeze(log_pis_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()
|
||||
|
||||
batch_size = tree.flatten(obs)[0].shape[0]
|
||||
if batch_size == policy.config["train_batch_size"]:
|
||||
policy.alpha_optim.zero_grad()
|
||||
alpha_loss.backward()
|
||||
policy.alpha_optim.step()
|
||||
|
||||
# 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.detach() * log_pis_t - bc_logp).mean()
|
||||
# actor_loss = -bc_logp.mean()
|
||||
|
||||
if batch_size == policy.config["train_batch_size"]:
|
||||
policy.actor_optim.zero_grad()
|
||||
actor_loss.backward(retain_graph=True)
|
||||
policy.actor_optim.step()
|
||||
|
||||
# Critic Loss (Standard SAC Critic L2 Loss + CQL Entropy Loss)
|
||||
# SAC Loss:
|
||||
# Q-values for the batched actions.
|
||||
action_dist_inputs_tp1, _ = model.get_action_model_outputs(model_out_tp1)
|
||||
action_dist_tp1 = action_dist_class(action_dist_inputs_tp1, model)
|
||||
policy_tp1, _ = action_dist_tp1.sample_logp()
|
||||
|
||||
q_t, _ = model.get_q_values(model_out_t, train_batch[SampleBatch.ACTIONS])
|
||||
q_t_selected = 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_selected = torch.squeeze(twin_q_t, dim=-1)
|
||||
|
||||
# Target q network evaluation.
|
||||
q_tp1, _ = target_model.get_q_values(target_model_out_tp1, policy_tp1)
|
||||
if twin_q:
|
||||
twin_q_tp1, _ = 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_best = torch.squeeze(input=q_tp1, dim=-1)
|
||||
q_tp1_best_masked = (1.0 - terminals.float()) * q_tp1_best
|
||||
|
||||
# compute RHS of bellman equation
|
||||
q_t_target = (
|
||||
rewards + (discount ** policy.config["n_step"]) * q_tp1_best_masked
|
||||
).detach()
|
||||
|
||||
# Compute the TD-error (potentially clipped), for priority replay buffer
|
||||
base_td_error = torch.abs(q_t_selected - q_t_target)
|
||||
if twin_q:
|
||||
twin_td_error = torch.abs(twin_q_t_selected - q_t_target)
|
||||
td_error = 0.5 * (base_td_error + twin_td_error)
|
||||
else:
|
||||
td_error = base_td_error
|
||||
|
||||
critic_loss_1 = nn.functional.mse_loss(q_t_selected, q_t_target)
|
||||
if twin_q:
|
||||
critic_loss_2 = nn.functional.mse_loss(twin_q_t_selected, 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
|
||||
),
|
||||
policy.device,
|
||||
)
|
||||
curr_actions, curr_logp = policy_actions_repeat(
|
||||
model, action_dist_class, model_out_t, num_actions
|
||||
)
|
||||
next_actions, next_logp = policy_actions_repeat(
|
||||
model, action_dist_class, model_out_tp1, num_actions
|
||||
)
|
||||
|
||||
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_qf1_loss]
|
||||
if twin_q:
|
||||
cql_loss.append(min_qf2_loss)
|
||||
|
||||
critic_loss = [critic_loss_1 + min_qf1_loss]
|
||||
if twin_q:
|
||||
critic_loss.append(critic_loss_2 + min_qf2_loss)
|
||||
|
||||
if batch_size == policy.config["train_batch_size"]:
|
||||
policy.critic_optims[0].zero_grad()
|
||||
critic_loss[0].backward(retain_graph=True)
|
||||
policy.critic_optims[0].step()
|
||||
|
||||
if twin_q:
|
||||
policy.critic_optims[1].zero_grad()
|
||||
critic_loss[1].backward(retain_graph=False)
|
||||
policy.critic_optims[1].step()
|
||||
|
||||
# Store values for stats function in model (tower), such that for
|
||||
# multi-GPU, we do not override them during the parallel loss phase.
|
||||
# SAC stats.
|
||||
model.tower_stats["q_t"] = q_t_selected
|
||||
model.tower_stats["policy_t"] = policy_t
|
||||
model.tower_stats["log_pis_t"] = log_pis_t
|
||||
model.tower_stats["actor_loss"] = actor_loss
|
||||
model.tower_stats["critic_loss"] = critic_loss
|
||||
model.tower_stats["alpha_loss"] = alpha_loss
|
||||
model.tower_stats["log_alpha_value"] = model.log_alpha
|
||||
model.tower_stats["alpha_value"] = alpha
|
||||
model.tower_stats["target_entropy"] = model.target_entropy
|
||||
# CQL stats.
|
||||
model.tower_stats["cql_loss"] = cql_loss
|
||||
|
||||
# TD-error tensor in final stats
|
||||
# will be concatenated and retrieved for each individual batch item.
|
||||
model.tower_stats["td_error"] = td_error
|
||||
|
||||
if use_lagrange:
|
||||
model.tower_stats["log_alpha_prime_value"] = model.log_alpha_prime[0]
|
||||
model.tower_stats["alpha_prime_value"] = alpha_prime
|
||||
model.tower_stats["alpha_prime_loss"] = alpha_prime_loss
|
||||
|
||||
if batch_size == policy.config["train_batch_size"]:
|
||||
policy.alpha_prime_optim.zero_grad()
|
||||
alpha_prime_loss.backward()
|
||||
policy.alpha_prime_optim.step()
|
||||
|
||||
# Return all loss terms corresponding to our optimizers.
|
||||
return tuple(
|
||||
[actor_loss]
|
||||
+ critic_loss
|
||||
+ [alpha_loss]
|
||||
+ ([alpha_prime_loss] if use_lagrange else [])
|
||||
)
|
||||
|
||||
|
||||
def cql_stats(policy: Policy, train_batch: SampleBatch) -> Dict[str, TensorType]:
|
||||
# Get SAC loss stats.
|
||||
stats_dict = stats(policy, train_batch)
|
||||
|
||||
# Add CQL loss stats to the dict.
|
||||
stats_dict["cql_loss"] = torch.mean(
|
||||
torch.stack(*policy.get_tower_stats("cql_loss"))
|
||||
)
|
||||
|
||||
if policy.config["lagrangian"]:
|
||||
stats_dict["log_alpha_prime_value"] = torch.mean(
|
||||
torch.stack(policy.get_tower_stats("log_alpha_prime_value"))
|
||||
)
|
||||
stats_dict["alpha_prime_value"] = torch.mean(
|
||||
torch.stack(policy.get_tower_stats("alpha_prime_value"))
|
||||
)
|
||||
stats_dict["alpha_prime_loss"] = torch.mean(
|
||||
torch.stack(policy.get_tower_stats("alpha_prime_loss"))
|
||||
)
|
||||
return stats_dict
|
||||
|
||||
|
||||
def cql_optimizer_fn(
|
||||
policy: Policy, config: AlgorithmConfigDict
|
||||
) -> 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: AlgorithmConfigDict,
|
||||
) -> 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)
|
||||
|
||||
|
||||
def compute_gradients_fn(policy, postprocessed_batch):
|
||||
batches = [policy._lazy_tensor_dict(postprocessed_batch)]
|
||||
model = policy.model
|
||||
policy._loss(policy, model, policy.dist_class, batches[0])
|
||||
stats = {LEARNER_STATS_KEY: policy._convert_to_numpy(cql_stats(policy, batches[0]))}
|
||||
return [None, stats]
|
||||
|
||||
|
||||
def apply_gradients_fn(policy, gradients):
|
||||
return
|
||||
|
||||
|
||||
# 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.algorithms.cql.cql.CQLConfig(),
|
||||
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,
|
||||
extra_learn_fetches_fn=concat_multi_gpu_td_errors,
|
||||
mixins=[TargetNetworkMixin, ComputeTDErrorMixin],
|
||||
action_distribution_fn=action_distribution_fn,
|
||||
compute_gradients_fn=compute_gradients_fn,
|
||||
apply_gradients_fn=apply_gradients_fn,
|
||||
)
|
||||
@@ -0,0 +1,128 @@
|
||||
import os
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import ray
|
||||
from ray.rllib.algorithms import cql
|
||||
from ray.rllib.utils.framework import try_import_torch
|
||||
from ray.rllib.utils.metrics import (
|
||||
ENV_RUNNER_RESULTS,
|
||||
EPISODE_RETURN_MEAN,
|
||||
EVALUATION_RESULTS,
|
||||
)
|
||||
from ray.rllib.utils.test_utils import check_compute_single_action, check_train_results
|
||||
|
||||
torch, _ = try_import_torch()
|
||||
|
||||
|
||||
class TestCQL(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
ray.init()
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
ray.shutdown()
|
||||
|
||||
def test_cql_compilation(self):
|
||||
"""Test whether CQL can be built with all frameworks."""
|
||||
|
||||
# Learns from a historic-data file.
|
||||
# To generate this data, first run:
|
||||
# $ ./train.py --run=SAC --env=Pendulum-v1 \
|
||||
# --stop='{"timesteps_total": 50000}' \
|
||||
# --config='{"output": "/tmp/out"}'
|
||||
rllib_dir = Path(__file__).parent.parent.parent.parent
|
||||
print("rllib dir={}".format(rllib_dir))
|
||||
data_file = os.path.join(rllib_dir, "offline/tests/data/pendulum/small.json")
|
||||
print("data_file={} exists={}".format(data_file, os.path.isfile(data_file)))
|
||||
|
||||
config = (
|
||||
cql.CQLConfig()
|
||||
.api_stack(
|
||||
enable_rl_module_and_learner=False,
|
||||
enable_env_runner_and_connector_v2=False,
|
||||
)
|
||||
.environment(
|
||||
env="Pendulum-v1",
|
||||
)
|
||||
.offline_data(
|
||||
input_=data_file,
|
||||
# In the files, we use here for testing, actions have already
|
||||
# been normalized.
|
||||
# This is usually the case when the file was generated by another
|
||||
# RLlib algorithm (e.g. PPO or SAC).
|
||||
actions_in_input_normalized=False,
|
||||
)
|
||||
.training(
|
||||
clip_actions=False,
|
||||
train_batch_size=2000,
|
||||
twin_q=True,
|
||||
num_steps_sampled_before_learning_starts=0,
|
||||
bc_iters=2,
|
||||
)
|
||||
.evaluation(
|
||||
evaluation_interval=2,
|
||||
evaluation_duration=10,
|
||||
evaluation_config=cql.CQLConfig.overrides(input_="sampler"),
|
||||
evaluation_parallel_to_training=False,
|
||||
evaluation_num_env_runners=2,
|
||||
)
|
||||
.env_runners(num_env_runners=0)
|
||||
.reporting(min_time_s_per_iteration=0)
|
||||
)
|
||||
num_iterations = 4
|
||||
|
||||
algo = config.build()
|
||||
for i in range(num_iterations):
|
||||
results = algo.train()
|
||||
check_train_results(results)
|
||||
print(results)
|
||||
eval_results = results.get(EVALUATION_RESULTS)
|
||||
if eval_results:
|
||||
print(
|
||||
f"iter={algo.iteration} "
|
||||
f"R={eval_results[ENV_RUNNER_RESULTS][EPISODE_RETURN_MEAN]}"
|
||||
)
|
||||
check_compute_single_action(algo)
|
||||
|
||||
# Get policy and model.
|
||||
pol = algo.get_policy()
|
||||
cql_model = pol.model
|
||||
|
||||
# Example on how to do evaluation on the trained Algorithm
|
||||
# using the data from CQL's global replay buffer.
|
||||
# Get a sample (MultiAgentBatch).
|
||||
|
||||
batch = algo.env_runner.input_reader.next()
|
||||
multi_agent_batch = batch.as_multi_agent()
|
||||
# All experiences have been buffered for `default_policy`
|
||||
batch = multi_agent_batch.policy_batches["default_policy"]
|
||||
|
||||
obs = torch.from_numpy(batch["obs"])
|
||||
|
||||
# Pass the observations through our model to get the
|
||||
# features, which then to pass through the Q-head.
|
||||
model_out, _ = cql_model({"obs": obs})
|
||||
# The estimated Q-values from the (historic) actions in the batch.
|
||||
q_values_old = cql_model.get_q_values(
|
||||
model_out, torch.from_numpy(batch["actions"])
|
||||
)
|
||||
|
||||
# The estimated Q-values for the new actions computed
|
||||
# by our policy.
|
||||
actions_new = pol.compute_actions_from_input_dict({"obs": obs})[0]
|
||||
q_values_new = cql_model.get_q_values(model_out, torch.from_numpy(actions_new))
|
||||
|
||||
print(f"Q-val batch={q_values_old}")
|
||||
print(f"Q-val policy={q_values_new}")
|
||||
|
||||
algo.stop()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,280 @@
|
||||
from typing import Dict
|
||||
|
||||
from ray.rllib.algorithms.cql.cql import CQLConfig
|
||||
from ray.rllib.algorithms.sac.sac_learner import (
|
||||
LOGPS_KEY,
|
||||
QF_LOSS_KEY,
|
||||
QF_MAX_KEY,
|
||||
QF_MEAN_KEY,
|
||||
QF_MIN_KEY,
|
||||
QF_PREDS,
|
||||
QF_TWIN_LOSS_KEY,
|
||||
QF_TWIN_PREDS,
|
||||
TD_ERROR_MEAN_KEY,
|
||||
)
|
||||
from ray.rllib.algorithms.sac.torch.sac_torch_learner import SACTorchLearner
|
||||
from ray.rllib.core.columns import Columns
|
||||
from ray.rllib.core.learner.learner import (
|
||||
POLICY_LOSS_KEY,
|
||||
)
|
||||
from ray.rllib.utils.annotations import override
|
||||
from ray.rllib.utils.framework import try_import_torch
|
||||
from ray.rllib.utils.metrics import ALL_MODULES
|
||||
from ray.rllib.utils.typing import ModuleID, ParamDict, TensorType
|
||||
from ray.tune.result import TRAINING_ITERATION
|
||||
|
||||
torch, nn = try_import_torch()
|
||||
|
||||
|
||||
class CQLTorchLearner(SACTorchLearner):
|
||||
@override(SACTorchLearner)
|
||||
def compute_loss_for_module(
|
||||
self,
|
||||
*,
|
||||
module_id: ModuleID,
|
||||
config: CQLConfig,
|
||||
batch: Dict,
|
||||
fwd_out: Dict[str, TensorType],
|
||||
) -> TensorType:
|
||||
|
||||
# TODO (simon, sven): Add upstream information pieces into this timesteps
|
||||
# call arg to Learner.update_...().
|
||||
self.metrics.log_value(
|
||||
(ALL_MODULES, TRAINING_ITERATION),
|
||||
1,
|
||||
reduce="sum",
|
||||
)
|
||||
# Get the train action distribution for the current policy and current state.
|
||||
# This is needed for the policy (actor) loss and the `alpha`` loss.
|
||||
action_dist_class = self.module[module_id].get_train_action_dist_cls()
|
||||
action_dist_curr = action_dist_class.from_logits(
|
||||
fwd_out[Columns.ACTION_DIST_INPUTS]
|
||||
)
|
||||
|
||||
# Optimize also the hyperparameter `alpha` by using the current policy
|
||||
# evaluated at the current state (from offline data). Note, in contrast
|
||||
# to the original SAC loss, here the `alpha` and actor losses are
|
||||
# calculated first.
|
||||
# TODO (simon): Check, why log(alpha) is used, prob. just better
|
||||
# to optimize and monotonic function. Original equation uses alpha.
|
||||
alpha_loss = -torch.mean(
|
||||
self.curr_log_alpha[module_id]
|
||||
* (fwd_out["logp_resampled"].detach() + self.target_entropy[module_id])
|
||||
)
|
||||
|
||||
# Get the current alpha.
|
||||
alpha = torch.exp(self.curr_log_alpha[module_id])
|
||||
# Start training with behavior cloning and turn to the classic Soft-Actor Critic
|
||||
# after `bc_iters` of training iterations.
|
||||
if (
|
||||
self.metrics.peek((ALL_MODULES, TRAINING_ITERATION), default=0)
|
||||
>= config.bc_iters
|
||||
):
|
||||
actor_loss = torch.mean(
|
||||
alpha.detach() * fwd_out["logp_resampled"] - fwd_out["q_curr"]
|
||||
)
|
||||
else:
|
||||
# Use log-probabilities of the current action distribution to clone
|
||||
# the behavior policy (selected actions in data) in the first `bc_iters`
|
||||
# training iterations.
|
||||
bc_logps_curr = action_dist_curr.logp(batch[Columns.ACTIONS])
|
||||
actor_loss = torch.mean(
|
||||
alpha.detach() * fwd_out["logp_resampled"] - bc_logps_curr
|
||||
)
|
||||
|
||||
# The critic loss is composed of the standard SAC Critic L2 loss and the
|
||||
# CQL entropy loss.
|
||||
|
||||
# Get the Q-values for the actually selected actions in the offline data.
|
||||
# In the critic loss we use these as predictions.
|
||||
q_selected = fwd_out[QF_PREDS]
|
||||
if config.twin_q:
|
||||
q_twin_selected = fwd_out[QF_TWIN_PREDS]
|
||||
|
||||
if not config.deterministic_backup:
|
||||
q_next = (
|
||||
fwd_out["q_target_next"]
|
||||
- alpha.detach() * fwd_out["logp_next_resampled"]
|
||||
)
|
||||
else:
|
||||
q_next = fwd_out["q_target_next"]
|
||||
|
||||
# Now mask all Q-values with terminating next states in the targets.
|
||||
q_next_masked = (1.0 - batch[Columns.TERMINATEDS].float()) * q_next
|
||||
|
||||
# Compute the right hand side of the Bellman equation. Detach this node
|
||||
# from the computation graph as we do not want to backpropagate through
|
||||
# the target network when optimizing the Q loss.
|
||||
# TODO (simon, sven): Kumar et al. (2020) use here also a reward scaler.
|
||||
q_selected_target = (
|
||||
# TODO (simon): Add an `n_step` option to the `AddNextObsToBatch` connector.
|
||||
batch[Columns.REWARDS]
|
||||
# TODO (simon): Implement n_step.
|
||||
+ (config.gamma) * q_next_masked
|
||||
).detach()
|
||||
|
||||
# Calculate the TD error.
|
||||
td_error = torch.abs(q_selected - q_selected_target)
|
||||
# Calculate a TD-error for twin-Q values, if needed.
|
||||
if config.twin_q:
|
||||
td_error += torch.abs(q_twin_selected - q_selected_target)
|
||||
# Rescale the TD error
|
||||
td_error *= 0.5
|
||||
|
||||
# MSBE loss for the critic(s) (i.e. Q, see eqs. (7-8) Haarnoja et al. (2018)).
|
||||
# Note, this needs a sample from the current policy given the next state.
|
||||
# Note further, we could also use here the Huber loss instead of the MSE.
|
||||
# TODO (simon): Add the huber loss as an alternative (SAC uses it).
|
||||
sac_critic_loss = torch.nn.MSELoss(reduction="mean")(
|
||||
q_selected,
|
||||
q_selected_target,
|
||||
)
|
||||
if config.twin_q:
|
||||
sac_critic_twin_loss = torch.nn.MSELoss(reduction="mean")(
|
||||
q_twin_selected,
|
||||
q_selected_target,
|
||||
)
|
||||
|
||||
# Now calculate the CQL loss (we use the entropy version of the CQL algorithm).
|
||||
# Note, the entropy version performs best in shown experiments.
|
||||
|
||||
# Compute the log-probabilities for the random actions (note, we generate random
|
||||
# actions (from the mu distribution as named in Kumar et al. (2020))).
|
||||
# Note, all actions, action log-probabilities and Q-values are already computed
|
||||
# by the module's `_forward_train` method.
|
||||
# TODO (simon): This is the density for a discrete uniform, however, actions
|
||||
# come from a continuous one. So actually this density should use (1/(high-low))
|
||||
# instead of (1/2).
|
||||
random_density = torch.log(
|
||||
torch.pow(
|
||||
0.5,
|
||||
torch.tensor(
|
||||
fwd_out["actions_curr_repeat"].shape[-1],
|
||||
device=fwd_out["actions_curr_repeat"].device,
|
||||
),
|
||||
)
|
||||
)
|
||||
# Merge all Q-values and subtract the log-probabilities (note, we use the
|
||||
# entropy version of CQL).
|
||||
q_repeat = torch.cat(
|
||||
[
|
||||
fwd_out["q_rand_repeat"] - random_density,
|
||||
fwd_out["q_next_repeat"] - fwd_out["logps_next_repeat"].detach(),
|
||||
fwd_out["q_curr_repeat"] - fwd_out["logps_curr_repeat"].detach(),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
cql_loss = (
|
||||
torch.logsumexp(q_repeat / config.temperature, dim=1).mean()
|
||||
* config.min_q_weight
|
||||
* config.temperature
|
||||
)
|
||||
cql_loss -= q_selected.mean() * config.min_q_weight
|
||||
# Add the CQL loss term to the SAC loss term.
|
||||
critic_loss = sac_critic_loss + cql_loss
|
||||
|
||||
# If a twin Q-value function is implemented calculated its CQL loss.
|
||||
if config.twin_q:
|
||||
q_twin_repeat = torch.cat(
|
||||
[
|
||||
fwd_out["q_twin_rand_repeat"] - random_density,
|
||||
fwd_out["q_twin_next_repeat"]
|
||||
- fwd_out["logps_next_repeat"].detach(),
|
||||
fwd_out["q_twin_curr_repeat"]
|
||||
- fwd_out["logps_curr_repeat"].detach(),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
cql_twin_loss = (
|
||||
torch.logsumexp(q_twin_repeat / config.temperature, dim=1).mean()
|
||||
* config.min_q_weight
|
||||
* config.temperature
|
||||
)
|
||||
cql_twin_loss -= q_twin_selected.mean() * config.min_q_weight
|
||||
# Add the CQL loss term to the SAC loss term.
|
||||
critic_twin_loss = sac_critic_twin_loss + cql_twin_loss
|
||||
|
||||
# TODO (simon): Check, if we need to implement here also a Lagrangian
|
||||
# loss.
|
||||
|
||||
total_loss = actor_loss + critic_loss + alpha_loss
|
||||
|
||||
# Add the twin critic loss to the total loss, if needed.
|
||||
if config.twin_q:
|
||||
# Reweigh the critic loss terms in the total loss.
|
||||
total_loss += 0.5 * critic_twin_loss - 0.5 * critic_loss
|
||||
|
||||
# Log important loss stats (reduce=mean (default), but with window=1
|
||||
# in order to keep them history free).
|
||||
self.metrics.log_dict(
|
||||
{
|
||||
POLICY_LOSS_KEY: actor_loss,
|
||||
QF_LOSS_KEY: critic_loss,
|
||||
# TODO (simon): Add these keys to SAC Learner.
|
||||
"cql_loss": cql_loss,
|
||||
"alpha_loss": alpha_loss,
|
||||
"alpha_value": alpha[0],
|
||||
"log_alpha_value": torch.log(alpha)[0],
|
||||
"target_entropy": self.target_entropy[module_id],
|
||||
LOGPS_KEY: torch.mean(
|
||||
fwd_out["logp_resampled"]
|
||||
), # torch.mean(logps_curr),
|
||||
QF_MEAN_KEY: torch.mean(fwd_out["q_curr_repeat"]),
|
||||
QF_MAX_KEY: torch.max(fwd_out["q_curr_repeat"]),
|
||||
QF_MIN_KEY: torch.min(fwd_out["q_curr_repeat"]),
|
||||
TD_ERROR_MEAN_KEY: torch.mean(td_error),
|
||||
},
|
||||
key=module_id,
|
||||
window=1, # <- single items (should not be mean/ema-reduced over time).
|
||||
)
|
||||
self._temp_losses[(module_id, POLICY_LOSS_KEY)] = actor_loss
|
||||
self._temp_losses[(module_id, QF_LOSS_KEY)] = critic_loss
|
||||
self._temp_losses[(module_id, "alpha_loss")] = alpha_loss
|
||||
|
||||
# TODO (simon): Add loss keys for langrangian, if needed.
|
||||
# TODO (simon): Add only here then the Langrange parameter optimization.
|
||||
if config.twin_q:
|
||||
self.metrics.log_value(
|
||||
key=(module_id, QF_TWIN_LOSS_KEY),
|
||||
value=critic_twin_loss,
|
||||
window=1, # <- single items (should not be mean/ema-reduced over time).
|
||||
)
|
||||
self._temp_losses[(module_id, QF_TWIN_LOSS_KEY)] = critic_twin_loss
|
||||
|
||||
# Return the total loss.
|
||||
return total_loss
|
||||
|
||||
@override(SACTorchLearner)
|
||||
def compute_gradients(
|
||||
self, loss_per_module: Dict[ModuleID, TensorType], **kwargs
|
||||
) -> ParamDict:
|
||||
|
||||
grads = {}
|
||||
for module_id in set(loss_per_module.keys()) - {ALL_MODULES}:
|
||||
# Loop through optimizers registered for this module.
|
||||
for optim_name, optim in self.get_optimizers_for_module(module_id):
|
||||
# Zero the gradients. Note, we need to reset the gradients b/c
|
||||
# each component for a module operates on the same graph.
|
||||
optim.zero_grad(set_to_none=True)
|
||||
|
||||
# Compute the gradients for the component and module.
|
||||
loss_tensor = self._temp_losses.pop((module_id, optim_name + "_loss"))
|
||||
loss_tensor.backward(
|
||||
retain_graph=False if optim_name in ["policy", "alpha"] else True
|
||||
)
|
||||
# Store the gradients for the component and module.
|
||||
# TODO (simon): Check another time the graph for overlapping
|
||||
# gradients.
|
||||
grads.update(
|
||||
{
|
||||
pid: grads[pid] + p.grad.clone()
|
||||
if pid in grads
|
||||
else p.grad.clone()
|
||||
for pid, p in self.filter_param_dict_for_optimizer(
|
||||
self._params, optim
|
||||
).items()
|
||||
}
|
||||
)
|
||||
|
||||
assert not self._temp_losses
|
||||
return grads
|
||||
@@ -0,0 +1,207 @@
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import tree
|
||||
|
||||
from ray.rllib.algorithms.sac.sac_catalog import SACCatalog
|
||||
from ray.rllib.algorithms.sac.sac_learner import (
|
||||
QF_PREDS,
|
||||
QF_TWIN_PREDS,
|
||||
)
|
||||
from ray.rllib.algorithms.sac.torch.default_sac_torch_rl_module import (
|
||||
DefaultSACTorchRLModule,
|
||||
)
|
||||
from ray.rllib.core.columns import Columns
|
||||
from ray.rllib.core.models.base import ENCODER_OUT
|
||||
from ray.rllib.utils.annotations import override
|
||||
from ray.rllib.utils.framework import try_import_torch
|
||||
from ray.rllib.utils.typing import TensorType
|
||||
|
||||
torch, nn = try_import_torch()
|
||||
|
||||
|
||||
class DefaultCQLTorchRLModule(DefaultSACTorchRLModule):
|
||||
def __init__(self, *args, **kwargs):
|
||||
catalog_class = kwargs.pop("catalog_class", None)
|
||||
if catalog_class is None:
|
||||
catalog_class = SACCatalog
|
||||
super().__init__(*args, **kwargs, catalog_class=catalog_class)
|
||||
|
||||
@override(DefaultSACTorchRLModule)
|
||||
def _forward_train(self, batch: Dict) -> Dict[str, Any]:
|
||||
# Call the super method.
|
||||
fwd_out = super()._forward_train(batch)
|
||||
|
||||
# Make sure we perform a "straight-through gradient" pass here,
|
||||
# ignoring the gradients of the q-net, however, still recording
|
||||
# the gradients of the policy net (which was used to rsample the actions used
|
||||
# here). This is different from doing `.detach()` or `with torch.no_grads()`,
|
||||
# as these two methds would fully block all gradient recordings, including
|
||||
# the needed policy ones.
|
||||
all_params = list(self.pi_encoder.parameters()) + list(self.pi.parameters())
|
||||
# if self.twin_q:
|
||||
# all_params += list(self.qf_twin.parameters()) + list(
|
||||
# self.qf_twin_encoder.parameters()
|
||||
# )
|
||||
|
||||
for param in all_params:
|
||||
param.requires_grad = False
|
||||
|
||||
# Compute the repeated actions, action log-probabilites and Q-values for all
|
||||
# observations.
|
||||
# First for the random actions (from the mu-distribution as named by Kumar et
|
||||
# al. (2020)).
|
||||
low = torch.tensor(
|
||||
self.action_space.low,
|
||||
device=fwd_out[QF_PREDS].device,
|
||||
)
|
||||
high = torch.tensor(
|
||||
self.action_space.high,
|
||||
device=fwd_out[QF_PREDS].device,
|
||||
)
|
||||
num_samples = batch[Columns.ACTIONS].shape[0] * self.model_config["num_actions"]
|
||||
actions_rand_repeat = low + (high - low) * torch.rand(
|
||||
(num_samples, low.shape[0]), device=fwd_out[QF_PREDS].device
|
||||
)
|
||||
|
||||
# First for the random actions (from the mu-distribution as named in Kumar
|
||||
# et al. (2020)) using repeated observations.
|
||||
rand_repeat_out = self._repeat_actions(batch[Columns.OBS], actions_rand_repeat)
|
||||
(fwd_out["actions_rand_repeat"], fwd_out["q_rand_repeat"]) = (
|
||||
rand_repeat_out[Columns.ACTIONS],
|
||||
rand_repeat_out[QF_PREDS],
|
||||
)
|
||||
# Sample current and next actions (from the pi distribution as named in Kumar
|
||||
# et al. (2020)) using repeated observations
|
||||
# Second for the current observations and the current action distribution.
|
||||
curr_repeat_out = self._repeat_actions(batch[Columns.OBS])
|
||||
(
|
||||
fwd_out["actions_curr_repeat"],
|
||||
fwd_out["logps_curr_repeat"],
|
||||
fwd_out["q_curr_repeat"],
|
||||
) = (
|
||||
curr_repeat_out[Columns.ACTIONS],
|
||||
curr_repeat_out[Columns.ACTION_LOGP],
|
||||
curr_repeat_out[QF_PREDS],
|
||||
)
|
||||
# Then, for the next observations and the current action distribution.
|
||||
next_repeat_out = self._repeat_actions(batch[Columns.NEXT_OBS])
|
||||
(
|
||||
fwd_out["actions_next_repeat"],
|
||||
fwd_out["logps_next_repeat"],
|
||||
fwd_out["q_next_repeat"],
|
||||
) = (
|
||||
next_repeat_out[Columns.ACTIONS],
|
||||
next_repeat_out[Columns.ACTION_LOGP],
|
||||
next_repeat_out[QF_PREDS],
|
||||
)
|
||||
if self.twin_q:
|
||||
# First for the random actions from the mu-distribution.
|
||||
fwd_out["q_twin_rand_repeat"] = rand_repeat_out[QF_TWIN_PREDS]
|
||||
# Second for the current observations and the current action distribution.
|
||||
fwd_out["q_twin_curr_repeat"] = curr_repeat_out[QF_TWIN_PREDS]
|
||||
# Then, for the next observations and the current action distribution.
|
||||
fwd_out["q_twin_next_repeat"] = next_repeat_out[QF_TWIN_PREDS]
|
||||
# Reset the gradient requirements for all Q-function parameters.
|
||||
for param in all_params:
|
||||
param.requires_grad = True
|
||||
|
||||
return fwd_out
|
||||
|
||||
def _repeat_tensor(self, tensor: TensorType, repeat: int) -> TensorType:
|
||||
"""Generates a repeated version of a tensor.
|
||||
|
||||
The repetition is done similar `np.repeat` and repeats each value
|
||||
instead of the complete vector.
|
||||
|
||||
Args:
|
||||
tensor: The tensor to be repeated.
|
||||
repeat: How often each value in the tensor should be repeated.
|
||||
|
||||
Returns:
|
||||
A tensor holding `repeat` repeated values of the input `tensor`
|
||||
"""
|
||||
# Insert the new dimension at axis 1 into the tensor.
|
||||
t_repeat = tensor.unsqueeze(1)
|
||||
# Repeat the tensor along the new dimension.
|
||||
t_repeat = torch.repeat_interleave(t_repeat, repeat, dim=1)
|
||||
# Stack the repeated values into the batch dimension.
|
||||
t_repeat = t_repeat.view(-1, *tensor.shape[1:])
|
||||
# Return the repeated tensor.
|
||||
return t_repeat
|
||||
|
||||
def _repeat_actions(
|
||||
self, obs: TensorType, actions: Optional[TensorType] = None
|
||||
) -> Dict[str, TensorType]:
|
||||
"""Generated actions and Q-values for repeated observations.
|
||||
|
||||
The `self.model_config["num_actions"]` define a multiplier
|
||||
used for generating `num_actions` as many actions as the batch size.
|
||||
Observations are repeated and then a model forward pass is made.
|
||||
|
||||
Args:
|
||||
obs: A batched observation tensor.
|
||||
actions: An optional batched actions tensor.
|
||||
|
||||
Returns:
|
||||
A dictionary holding the (sampled or passed-in actions), the log
|
||||
probabilities (of sampled actions), the Q-values and if available
|
||||
the twin-Q values.
|
||||
"""
|
||||
output = {}
|
||||
# Receive the batch size.
|
||||
batch_size = obs.shape[0]
|
||||
# Receive the number of action to sample.
|
||||
num_actions = self.model_config["num_actions"]
|
||||
# Repeat the observations `num_actions` times.
|
||||
obs_repeat = tree.map_structure(
|
||||
lambda t: self._repeat_tensor(t, num_actions), obs
|
||||
)
|
||||
# Generate a batch for the forward pass.
|
||||
temp_batch = {Columns.OBS: obs_repeat}
|
||||
if actions is None:
|
||||
# TODO (simon): Run the forward pass in inference mode.
|
||||
# Compute the action logits.
|
||||
pi_encoder_outs = self.pi_encoder(temp_batch)
|
||||
action_logits = self.pi(pi_encoder_outs[ENCODER_OUT])
|
||||
# Generate the squashed Gaussian from the model's logits.
|
||||
action_dist = self.get_train_action_dist_cls().from_logits(action_logits)
|
||||
# Sample the actions. Note, we want to make a backward pass through
|
||||
# these actions.
|
||||
output[Columns.ACTIONS] = action_dist.rsample()
|
||||
# Compute the action log-probabilities.
|
||||
output[Columns.ACTION_LOGP] = action_dist.logp(
|
||||
output[Columns.ACTIONS]
|
||||
).view(batch_size, num_actions, 1)
|
||||
else:
|
||||
output[Columns.ACTIONS] = actions
|
||||
|
||||
# Compute all Q-values.
|
||||
temp_batch.update(
|
||||
{
|
||||
Columns.ACTIONS: output[Columns.ACTIONS],
|
||||
}
|
||||
)
|
||||
output.update(
|
||||
{
|
||||
QF_PREDS: self._qf_forward_train_helper(
|
||||
temp_batch,
|
||||
self.qf_encoder,
|
||||
self.qf,
|
||||
).view(batch_size, num_actions, 1)
|
||||
}
|
||||
)
|
||||
# If we have a twin-Q network, compute its Q-values, too.
|
||||
if self.twin_q:
|
||||
output.update(
|
||||
{
|
||||
QF_TWIN_PREDS: self._qf_forward_train_helper(
|
||||
temp_batch,
|
||||
self.qf_twin_encoder,
|
||||
self.qf_twin,
|
||||
).view(batch_size, num_actions, 1)
|
||||
}
|
||||
)
|
||||
del temp_batch
|
||||
|
||||
# Return
|
||||
return output
|
||||
Reference in New Issue
Block a user