wehub-resource-sync
116 lines
3.9 KiB
Python
116 lines
3.9 KiB
Python
from typing import Dict
|
|
|
|
import numpy as np
|
|
|
|
from ray.rllib.algorithms.dqn.dqn_learner import DQNLearner
|
|
from ray.rllib.core.learner.learner import Learner
|
|
from ray.rllib.utils.annotations import override
|
|
from ray.rllib.utils.lambda_defaultdict import LambdaDefaultDict
|
|
from ray.rllib.utils.typing import ModuleID, TensorType
|
|
|
|
# Now, this is double defined: In `DefaultSACRLModule` and here. I would keep it here
|
|
# or push it into the `Learner` as these are recurring keys in RL.
|
|
LOGPS_KEY = "logps"
|
|
QF_LOSS_KEY = "qf_loss"
|
|
QF_MEAN_KEY = "qf_mean"
|
|
QF_MAX_KEY = "qf_max"
|
|
QF_MIN_KEY = "qf_min"
|
|
QF_PREDS = "qf_preds"
|
|
QF_TWIN_LOSS_KEY = "qf_twin_loss"
|
|
QF_TWIN_PREDS = "qf_twin_preds"
|
|
TD_ERROR_MEAN_KEY = "td_error_mean"
|
|
CRITIC_TARGET = "critic_target"
|
|
ACTION_DIST_INPUTS_NEXT = "action_dist_inputs_next"
|
|
QF_TARGET_NEXT = "q_target_next"
|
|
ACTION_PROBS_NEXT = "action_probs_next"
|
|
ACTION_LOG_PROBS_NEXT = "action_log_probs_next"
|
|
ACTION_PROBS = "action_probs"
|
|
ACTION_LOG_PROBS = "action_log_probs"
|
|
|
|
|
|
class SACLearner(DQNLearner):
|
|
@override(Learner)
|
|
def build(self) -> None:
|
|
# Store the current alpha in log form. We need it during optimization
|
|
# in log form.
|
|
self.curr_log_alpha: Dict[ModuleID, TensorType] = LambdaDefaultDict(
|
|
lambda module_id: self._get_tensor_variable(
|
|
# Note, we want to train the temperature parameter.
|
|
[
|
|
np.log(
|
|
self.config.get_config_for_module(module_id).initial_alpha
|
|
).astype(np.float32)
|
|
],
|
|
trainable=True,
|
|
)
|
|
)
|
|
|
|
# We need to call the `super()`'s `build()` method here to have the variables
|
|
# for the alpha already defined.
|
|
super().build()
|
|
|
|
self.target_entropy: Dict[ModuleID, TensorType] = LambdaDefaultDict(
|
|
lambda module_id: self._get_tensor_variable(
|
|
self._get_target_entropy(module_id)
|
|
)
|
|
)
|
|
|
|
@override(Learner)
|
|
def remove_module(self, module_id: ModuleID) -> None:
|
|
"""Removes the temperature and target entropy.
|
|
|
|
Note, this means that we also need to remove the corresponding
|
|
temperature optimizer.
|
|
"""
|
|
super().remove_module(module_id)
|
|
self.curr_log_alpha.pop(module_id, None)
|
|
self.target_entropy.pop(module_id, None)
|
|
|
|
@override(Learner)
|
|
def add_module(
|
|
self,
|
|
*,
|
|
module_id,
|
|
module_spec,
|
|
config_overrides=None,
|
|
new_should_module_be_updated=None
|
|
):
|
|
# First call `super`'s `add_module` method.
|
|
super().add_module(
|
|
module_id=module_id,
|
|
module_spec=module_spec,
|
|
config_overrides=config_overrides,
|
|
new_should_module_be_updated=new_should_module_be_updated,
|
|
)
|
|
# Now add the log alpha.
|
|
self.curr_log_alpha[module_id] = self._get_tensor_variable(
|
|
# Note, we want to train the temperature parameter.
|
|
[
|
|
np.log(
|
|
self.config.get_config_for_module(module_id).initial_alpha
|
|
).astype(np.float32)
|
|
],
|
|
trainable=True,
|
|
)
|
|
# Add also the target entropy for the new module.
|
|
self.target_entropy[module_id] = self._get_tensor_variable(
|
|
self._get_target_entropy(module_id)
|
|
)
|
|
|
|
def _get_target_entropy(self, module_id):
|
|
"""Returns the target entropy to use for the loss.
|
|
|
|
Args:
|
|
module_id: Module ID for which the target entropy should be
|
|
returned.
|
|
|
|
Returns:
|
|
Target entropy.
|
|
"""
|
|
target_entropy = self.config.get_config_for_module(module_id).target_entropy
|
|
if target_entropy is None or target_entropy == "auto":
|
|
target_entropy = -np.prod(
|
|
self._module_spec.module_specs[module_id].action_space.shape
|
|
)
|
|
return target_entropy
|