chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,7 @@
|
||||
from ray.rllib.core.learner.learner import Learner
|
||||
from ray.rllib.core.learner.learner_group import LearnerGroup
|
||||
|
||||
__all__ = [
|
||||
"Learner",
|
||||
"LearnerGroup",
|
||||
]
|
||||
@@ -0,0 +1,813 @@
|
||||
import abc
|
||||
import logging
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Collection,
|
||||
Dict,
|
||||
Iterable,
|
||||
Optional,
|
||||
Tuple,
|
||||
Union,
|
||||
)
|
||||
|
||||
import numpy
|
||||
|
||||
from ray.rllib.connectors.learner.learner_connector_pipeline import (
|
||||
LearnerConnectorPipeline,
|
||||
)
|
||||
from ray.rllib.core import ALL_MODULES, COMPONENT_METRICS_LOGGER
|
||||
from ray.rllib.core.learner.training_data import TrainingData
|
||||
from ray.rllib.core.rl_module.multi_rl_module import MultiRLModule
|
||||
from ray.rllib.core.rl_module.rl_module import RLModule
|
||||
from ray.rllib.policy.sample_batch import MultiAgentBatch, SampleBatch
|
||||
from ray.rllib.utils import unflatten_dict
|
||||
from ray.rllib.utils.annotations import (
|
||||
OverrideToImplementCustomLogic,
|
||||
OverrideToImplementCustomLogic_CallToSuperRecommended,
|
||||
override,
|
||||
)
|
||||
from ray.rllib.utils.checkpoints import Checkpointable
|
||||
from ray.rllib.utils.metrics import (
|
||||
DATASET_NUM_ITERS_TRAINED,
|
||||
DATASET_NUM_ITERS_TRAINED_LIFETIME,
|
||||
MODULE_TRAIN_BATCH_SIZE_MEAN,
|
||||
NUM_ENV_STEPS_TRAINED,
|
||||
NUM_ENV_STEPS_TRAINED_LIFETIME,
|
||||
NUM_MODULE_STEPS_TRAINED,
|
||||
NUM_MODULE_STEPS_TRAINED_LIFETIME,
|
||||
WEIGHTS_SEQ_NO,
|
||||
)
|
||||
from ray.rllib.utils.metrics.metrics_logger import MetricsLogger
|
||||
from ray.rllib.utils.minibatch_utils import (
|
||||
MiniBatchCyclicIterator,
|
||||
MiniBatchDummyIterator,
|
||||
MiniBatchRayDataIterator,
|
||||
)
|
||||
from ray.rllib.utils.typing import (
|
||||
DeviceType,
|
||||
ModuleID,
|
||||
NamedParamDict,
|
||||
ResultDict,
|
||||
StateDict,
|
||||
TensorType,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
|
||||
from ray.rllib.core.learner.differentiable_learner_config import (
|
||||
DifferentiableLearnerConfig,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DifferentiableLearner(Checkpointable):
|
||||
"""A differentiable `Learner` class enabling functional parameter updates.
|
||||
|
||||
This class is part of RLlib's Meta Learning API and provides a differentiable
|
||||
`Learner`, allowing higher-order updates within the meta-learning loop.
|
||||
|
||||
Unlike standard learners, this class operates on a provided `RLModule` reference
|
||||
instead of creating its own. Updated cloned module parameters are returned to the
|
||||
caller, facilitating further update processes.
|
||||
|
||||
The `update`, `_update`, `compute_gradients`, and `apply_gradients` methods
|
||||
require, in addition to other arguments, a dictionary of cloned `RLModule`
|
||||
parameters for functional updates.
|
||||
"""
|
||||
|
||||
# The framework an instance uses.
|
||||
framework_str: str
|
||||
|
||||
# The key for the total loss.
|
||||
TOTAL_LOSS_KEY: str = "total_loss"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
config: "AlgorithmConfig",
|
||||
learner_config: "DifferentiableLearnerConfig",
|
||||
module: Optional[RLModule] = None,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
|
||||
# An `AlgorithmConfig` used to access certain global configurations.
|
||||
self.config: "AlgorithmConfig" = config.copy(copy_frozen=False)
|
||||
# The specific configuration for the differentiable learner instance.
|
||||
self.learner_config: "DifferentiableLearnerConfig" = learner_config
|
||||
# The reference to the caller's module.
|
||||
self._module: Optional[MultiRLModule] = module
|
||||
# The reference to the caller's device.
|
||||
self._device: Optional[DeviceType] = None
|
||||
# A counter for functional weight updates.
|
||||
self._weights_seq_no: int = 0
|
||||
|
||||
# Whether self.build has already been called.
|
||||
self._is_built: bool = False
|
||||
|
||||
# The learner connector pipeline.
|
||||
self._learner_connector: Optional[LearnerConnectorPipeline] = None
|
||||
|
||||
# The Learner's own MetricsLogger to be used to log RLlib's built-in metrics or
|
||||
# custom user-defined ones (e.g. custom loss values). When returning from an
|
||||
# `update_from_...()` method call, the Learner will do a `self.metrics.reduce()`
|
||||
# and return the resulting (reduced) dict.
|
||||
self.metrics: MetricsLogger = MetricsLogger(
|
||||
stats_cls_lookup=config.stats_cls_lookup,
|
||||
root=False,
|
||||
)
|
||||
|
||||
# In case of offline learning and multiple learners, each learner receives a
|
||||
# repeatable iterator that iterates over a split of the streamed data.
|
||||
self.iterator: MiniBatchRayDataIterator = None
|
||||
|
||||
@OverrideToImplementCustomLogic_CallToSuperRecommended
|
||||
def build(self, device: Optional[DeviceType] = None) -> None:
|
||||
|
||||
if self._is_built:
|
||||
logger.debug("DifferentiableLearner already built. Skipping built.")
|
||||
|
||||
# If a device was passed, set the `DifferentiableLearner`'s device.
|
||||
if device:
|
||||
self._device = device
|
||||
|
||||
# TODO (simon): Move the `build_learner_connector` to the
|
||||
# `DifferentiableLearnerConfig`.
|
||||
self._learner_connector = self.learner_config.build_learner_connector(
|
||||
input_observation_space=None,
|
||||
input_action_space=None,
|
||||
device=self._device,
|
||||
)
|
||||
|
||||
# This instance is now ready for use.
|
||||
self._is_built = True
|
||||
|
||||
@OverrideToImplementCustomLogic
|
||||
@abc.abstractmethod
|
||||
def compute_gradients(
|
||||
self,
|
||||
loss_per_module: Dict[ModuleID, TensorType],
|
||||
params: Dict[ModuleID, NamedParamDict],
|
||||
**kwargs,
|
||||
) -> Dict[ModuleID, NamedParamDict]:
|
||||
"""Computes functional gradients based on the given losses.
|
||||
|
||||
Note that this method requires computing gradients functionally,
|
||||
without relying on an optimizer. If an optimizer is needed, a
|
||||
differentiable optimizer from a third-party package must be used.
|
||||
|
||||
Args:
|
||||
loss_per_module: Dict mapping module IDs to their individual total loss
|
||||
terms, computed by the individual `compute_loss_for_module()` calls.
|
||||
The overall total loss (sum of loss terms over all modules) is stored
|
||||
under `loss_per_module[ALL_MODULES]`.
|
||||
params: A dictionary containing cloned parameters for each module id.
|
||||
**kwargs: Forward compatibility kwargs.
|
||||
|
||||
Returns:
|
||||
The gradients in the same (dict) format as `params`.
|
||||
"""
|
||||
|
||||
@OverrideToImplementCustomLogic
|
||||
@abc.abstractmethod
|
||||
def apply_gradients(
|
||||
self,
|
||||
gradients: Dict[ModuleID, NamedParamDict],
|
||||
params: Dict[ModuleID, NamedParamDict],
|
||||
) -> Dict[ModuleID, NamedParamDict]:
|
||||
"""Applies given gradients functionally.
|
||||
|
||||
Note that this method requires functional parameter updates,
|
||||
meaning modifications must not be performed in-place (e.g., via an
|
||||
optimizer or directly within the `MultiRLModule`).
|
||||
|
||||
Args:
|
||||
gradients: A dictionary containing named gradients for each module id.
|
||||
params: A dictionary containing named parameters for each module id.
|
||||
|
||||
Returns:
|
||||
The updated parameters in the same (dict) format as `params`.
|
||||
"""
|
||||
|
||||
@OverrideToImplementCustomLogic
|
||||
def should_module_be_updated(self, module_id, multi_agent_batch=None):
|
||||
"""Returns whether a module should be updated or not based on `self.config`.
|
||||
|
||||
Args:
|
||||
module_id: The ModuleID that we want to query on whether this module
|
||||
should be updated or not.
|
||||
multi_agent_batch: An optional MultiAgentBatch to possibly provide further
|
||||
information on the decision on whether the RLModule should be updated
|
||||
or not.
|
||||
"""
|
||||
should_module_be_updated_fn = self.config.policies_to_train
|
||||
# If None, return True (by default, all modules should be updated).
|
||||
if should_module_be_updated_fn is None:
|
||||
return True
|
||||
# If collection given, return whether `module_id` is in that container.
|
||||
elif not callable(should_module_be_updated_fn):
|
||||
return module_id in set(should_module_be_updated_fn)
|
||||
|
||||
return should_module_be_updated_fn(module_id, multi_agent_batch)
|
||||
|
||||
@OverrideToImplementCustomLogic
|
||||
def compute_losses(
|
||||
self, *, fwd_out: Dict[str, Any], batch: Dict[str, Any]
|
||||
) -> Dict[str, Any]:
|
||||
"""Computes the loss(es) for the module being optimized.
|
||||
|
||||
This method must be overridden by MultiRLModule-specific Learners in order to
|
||||
define the specific loss computation logic. If the algorithm is single-agent,
|
||||
only `compute_loss_for_module()` should be overridden instead. If the algorithm
|
||||
uses independent multi-agent learning (default behavior for RLlib's multi-agent
|
||||
setups), also only `compute_loss_for_module()` should be overridden, but it will
|
||||
be called for each individual RLModule inside the MultiRLModule.
|
||||
For the functional update to work, no `forward` call should be made
|
||||
within this method, especially not a non-functional one. Instead, use
|
||||
the model outputs provided by `fwd_out`.
|
||||
|
||||
Args:
|
||||
fwd_out: Output from a functional call to the `forward_train()` method of the
|
||||
underlying MultiRLModule (`self.module`) during training
|
||||
(`self.update()`).
|
||||
batch: The train batch that was used to compute `fwd_out`.
|
||||
|
||||
Returns:
|
||||
A dictionary mapping module IDs to individual loss terms.
|
||||
"""
|
||||
loss_per_module = {}
|
||||
for module_id in fwd_out:
|
||||
loss = self.compute_loss_for_module(
|
||||
module_id=module_id,
|
||||
# TODO (simon): Check, if this should be provided per
|
||||
# `DifferentiableLearnerConfig`.
|
||||
config=self.config.get_config_for_module(module_id),
|
||||
batch=batch[module_id],
|
||||
fwd_out=fwd_out[module_id],
|
||||
)
|
||||
loss_per_module[module_id] = loss
|
||||
|
||||
return loss_per_module
|
||||
|
||||
@OverrideToImplementCustomLogic
|
||||
@abc.abstractmethod
|
||||
def compute_loss_for_module(
|
||||
self,
|
||||
*,
|
||||
module_id: ModuleID,
|
||||
config: "AlgorithmConfig",
|
||||
batch: Dict[str, Any],
|
||||
fwd_out: Dict[str, TensorType],
|
||||
) -> TensorType:
|
||||
"""Computes the loss for a single module.
|
||||
|
||||
Think of this as computing loss for a single agent. For multi-agent use-cases
|
||||
that require more complicated computation for loss, consider overriding the
|
||||
`compute_losses` method instead.
|
||||
|
||||
Args:
|
||||
module_id: The id of the module.
|
||||
config: The AlgorithmConfig specific to the given `module_id`.
|
||||
batch: The train batch for this particular module.
|
||||
fwd_out: The output of the forward pass for this particular module.
|
||||
|
||||
Returns:
|
||||
A single total loss tensor. If you have more than one optimizer on the
|
||||
provided `module_id` and would like to compute gradients separately using
|
||||
these different optimizers, simply add up the individual loss terms for
|
||||
each optimizer and return the sum. Also, for recording/logging any
|
||||
individual loss terms, you can use the `Learner.metrics.log_value(
|
||||
key=..., value=...)` or `DifferentiableLearner.metrics.log_dict()` APIs.
|
||||
See: :py:class:`~ray.rllib.utils.metrics.metrics_logger.MetricsLogger` for
|
||||
more information.
|
||||
"""
|
||||
|
||||
def update(
|
||||
self,
|
||||
params: Dict[ModuleID, NamedParamDict],
|
||||
training_data: TrainingData,
|
||||
*,
|
||||
_no_metrics_reduce: bool = False,
|
||||
**kwargs,
|
||||
) -> Tuple[Dict[ModuleID, NamedParamDict], ResultDict]:
|
||||
"""Make a functional update on provided parameters.
|
||||
|
||||
You can use this method to take more than one backward pass on the batch. All
|
||||
configuration parameters for the iteration loop are set within the
|
||||
`learner_config`. Note, the same configuration will be used for all module ids
|
||||
in `MultiRLModule`.
|
||||
|
||||
Args:
|
||||
params: A parameter dictionary holding named parameters for each module id.
|
||||
These parameters must be a clone of the module's original parameters to
|
||||
perform a functional update on them.
|
||||
training_data: A `TrainingData` instance containing the data or data iterators
|
||||
to be used in updating the given parameters in `param`.
|
||||
|
||||
Returns:
|
||||
The functionally updated parameters in the (dict) format they were passed in
|
||||
and a `ResultDict` object produced by a call to `self.metrics.reduce()`. The
|
||||
returned dict may be arbitrarily nested and must have `Stats` objects at
|
||||
all its leafs, allowing components further downstream (i.e. a user of this
|
||||
Learner) to further reduce these results (for example over n parallel
|
||||
Learners).
|
||||
"""
|
||||
self._check_is_built()
|
||||
|
||||
# TODO (simon): Implement a `before_gradient_based_update`, if necessary.
|
||||
# Call `before_gradient_based_update` to allow for non-gradient based
|
||||
# preparations-, logging-, and update logic to happen.
|
||||
# self.before_gradient_based_update(timesteps=timesteps or {})
|
||||
|
||||
training_data.validate()
|
||||
training_data.solve_refs()
|
||||
assert training_data.batches is None, "`training_data.batches` must be None!"
|
||||
|
||||
self._weights_seq_no += 1
|
||||
|
||||
batch_iter = self._create_iterator_if_necessary(
|
||||
training_data=training_data,
|
||||
num_total_minibatches=self.learner_config.num_total_minibatches,
|
||||
num_epochs=self.learner_config.num_epochs,
|
||||
minibatch_size=self.learner_config.minibatch_size,
|
||||
shuffle_batch_per_epoch=self.learner_config.shuffle_batch_per_epoch,
|
||||
)
|
||||
|
||||
# Perform the actual looping through the minibatches or the given data iterator.
|
||||
for iteration, tensor_minibatch in enumerate(batch_iter):
|
||||
# Check the MultiAgentBatch, whether our RLModule contains all ModuleIDs
|
||||
# found in this batch. If not, throw an error.
|
||||
unknown_module_ids = set(tensor_minibatch.policy_batches.keys()) - set(
|
||||
self.module.keys()
|
||||
)
|
||||
if unknown_module_ids:
|
||||
raise ValueError(
|
||||
f"Batch contains one or more ModuleIDs ({unknown_module_ids}) that "
|
||||
f"are not in this Learner!"
|
||||
)
|
||||
|
||||
# Make the actual in-graph/traced `_update` call. This should return
|
||||
# all tensor values (no numpy).
|
||||
fwd_out, loss_per_module, params, _ = self._update(
|
||||
# TODO (simon): Maybe filter ParamDict by module keys in batch.
|
||||
tensor_minibatch.policy_batches,
|
||||
params,
|
||||
)
|
||||
|
||||
# TODO (sven): Maybe move this into loop above to get metrics more accuratcely
|
||||
# cover the minibatch/epoch logic.
|
||||
# Log all timesteps (env, agent, modules) based on given episodes/batch.
|
||||
self._log_steps_trained_metrics(tensor_minibatch)
|
||||
|
||||
self._set_slicing_by_batch_id(tensor_minibatch, value=False)
|
||||
|
||||
if self.iterator:
|
||||
# Record the number of batches pulled from the dataset.
|
||||
self.metrics.log_value(
|
||||
(ALL_MODULES, DATASET_NUM_ITERS_TRAINED),
|
||||
iteration + 1,
|
||||
reduce="sum",
|
||||
)
|
||||
self.metrics.log_value(
|
||||
(ALL_MODULES, DATASET_NUM_ITERS_TRAINED_LIFETIME),
|
||||
iteration + 1,
|
||||
reduce="lifetime_sum",
|
||||
)
|
||||
# Log all individual RLModules' loss terms and its registered optimizers'
|
||||
# current learning rates.
|
||||
# Note: We do this only once for the last of the minibatch updates, b/c the
|
||||
# window is only 1 anyways.
|
||||
for mid, loss in loss_per_module.items():
|
||||
self.metrics.log_value(
|
||||
key=(mid, self.TOTAL_LOSS_KEY),
|
||||
value=loss,
|
||||
window=1,
|
||||
)
|
||||
|
||||
# Call `after_gradient_based_update` to allow for non-gradient based
|
||||
# cleanups-, logging-, and update logic to happen.
|
||||
# TODO (simon): Check, if this should stay here, when running multiple
|
||||
# gradient steps inside the iterator loop above (could be a complete epoch)
|
||||
# the target networks might need to be updated earlier.
|
||||
# self.after_gradient_based_update(timesteps=timesteps or {})
|
||||
|
||||
# Reduce results across all minibatch update steps.
|
||||
if not _no_metrics_reduce:
|
||||
return params, loss_per_module, self.metrics.reduce()
|
||||
else:
|
||||
return params, loss_per_module, {}
|
||||
|
||||
def _create_iterator_if_necessary(
|
||||
self,
|
||||
*,
|
||||
training_data: TrainingData,
|
||||
num_total_minibatches: int = 0,
|
||||
num_epochs: int = 1,
|
||||
minibatch_size: Optional[int] = None,
|
||||
shuffle_batch_per_epoch: bool = False,
|
||||
**kwargs,
|
||||
) -> Iterable:
|
||||
"""Provides a batch iterator."""
|
||||
|
||||
# Data iterator provided.
|
||||
if training_data.data_iterators:
|
||||
num_iters = kwargs.pop("num_iters", None)
|
||||
if num_iters is None:
|
||||
raise ValueError(
|
||||
"Learner.update(data_iterators=..) requires `num_iters` kwarg!"
|
||||
)
|
||||
|
||||
def _collate_fn(_batch: Dict[str, numpy.ndarray]) -> MultiAgentBatch:
|
||||
_batch = unflatten_dict(_batch)
|
||||
_batch = MultiAgentBatch(
|
||||
{
|
||||
module_id: SampleBatch(module_data)
|
||||
for module_id, module_data in _batch.items()
|
||||
},
|
||||
env_steps=sum(
|
||||
len(next(iter(module_data.values())))
|
||||
for module_data in _batch.values()
|
||||
),
|
||||
)
|
||||
_batch = self._convert_batch_type(_batch, to_device=False)
|
||||
return self._set_slicing_by_batch_id(_batch, value=True)
|
||||
|
||||
def _finalize_fn(batch: MultiAgentBatch) -> MultiAgentBatch:
|
||||
return self._convert_batch_type(batch, to_device=True, use_stream=True)
|
||||
|
||||
if not self.iterator:
|
||||
# This iterator holds a `ray.data.DataIterator` and manages it state.
|
||||
self.iterator = MiniBatchRayDataIterator(
|
||||
iterator=training_data.data_iterators[0],
|
||||
collate_fn=_collate_fn,
|
||||
finalize_fn=_finalize_fn,
|
||||
minibatch_size=minibatch_size,
|
||||
num_iters=num_iters,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
batch_iter = self.iterator
|
||||
else:
|
||||
batch = self._make_batch_if_necessary(training_data=training_data)
|
||||
assert batch is not None
|
||||
|
||||
# TODO: Move this into LearnerConnector pipeline?
|
||||
# Filter out those RLModules from the final train batch that should not be
|
||||
# updated.
|
||||
for module_id in list(batch.policy_batches.keys()):
|
||||
if not self.should_module_be_updated(module_id, batch):
|
||||
del batch.policy_batches[module_id]
|
||||
if not batch.policy_batches:
|
||||
return {}
|
||||
|
||||
batch = self._set_slicing_by_batch_id(batch, value=True)
|
||||
|
||||
if minibatch_size:
|
||||
batch_iter_cls = MiniBatchCyclicIterator
|
||||
elif num_epochs > 1:
|
||||
# `minibatch_size` was not set but `num_epochs` > 1.
|
||||
minibatch_size = batch.count
|
||||
# Note that there is no need to shuffle here, b/c we don't have
|
||||
# minibatches.
|
||||
batch_iter_cls = MiniBatchCyclicIterator
|
||||
else:
|
||||
# `minibatch_size` and `num_epochs` are not set by the user.
|
||||
batch_iter_cls = MiniBatchDummyIterator
|
||||
|
||||
batch_iter = batch_iter_cls(
|
||||
batch,
|
||||
num_epochs=num_epochs,
|
||||
minibatch_size=minibatch_size,
|
||||
shuffle_batch_per_epoch=shuffle_batch_per_epoch and (num_epochs > 1),
|
||||
num_total_minibatches=num_total_minibatches,
|
||||
)
|
||||
return batch_iter
|
||||
|
||||
@OverrideToImplementCustomLogic
|
||||
@abc.abstractmethod
|
||||
def _update(
|
||||
self,
|
||||
batch: Dict[str, Any],
|
||||
params: Dict[ModuleID, NamedParamDict],
|
||||
**kwargs,
|
||||
) -> Tuple[Any, Dict[ModuleID, NamedParamDict], Any, Any]:
|
||||
"""Contains all logic for an in-graph/traceable functional update step.
|
||||
|
||||
Framework specific subclasses must implement this method. This should include
|
||||
functional calls to the RLModule's `forward_train`, `compute_loss`,
|
||||
`compute_gradients`, `postprocess_gradients`, and `apply_gradients` methods
|
||||
and return a tuple with all the individual results as well as the functionally
|
||||
updated parameter dictionary.
|
||||
|
||||
Args:
|
||||
batch: The train batch already converted to a Dict mapping str to (possibly
|
||||
nested) tensors.
|
||||
kwargs: Forward compatibility kwargs.
|
||||
|
||||
Returns:
|
||||
A tuple consisting of:
|
||||
1) the output of a functional forward call to the RLModule using
|
||||
`params`,
|
||||
2) the functionally updated parameters in the (dict) format passed in,
|
||||
3) the `loss_per_module` dictionary mapping module IDs to individual
|
||||
loss tensors,
|
||||
4) a metrics dict mapping module IDs to metrics key/value pairs.
|
||||
|
||||
"""
|
||||
|
||||
@override(Checkpointable)
|
||||
def get_state(
|
||||
self,
|
||||
components: Optional[Union[str, Collection[str]]] = None,
|
||||
*,
|
||||
not_components: Optional[Union[str, Collection[str]]] = None,
|
||||
**kwargs,
|
||||
) -> StateDict:
|
||||
"""Gets the state of the `DifferentiableLearner` instance.
|
||||
|
||||
Note, because the `MultiRLModule` held by this class is only a reference
|
||||
it is not contained in the class' state.
|
||||
"""
|
||||
self._check_is_built()
|
||||
|
||||
state = {
|
||||
"should_module_be_updated": self.config.policies_to_train,
|
||||
WEIGHTS_SEQ_NO: self._weights_seq_no,
|
||||
}
|
||||
|
||||
if self._check_component(COMPONENT_METRICS_LOGGER, components, not_components):
|
||||
# TODO (sven): Make `MetricsLogger` a Checkpointable.
|
||||
state[COMPONENT_METRICS_LOGGER] = self.metrics.get_state()
|
||||
|
||||
return state
|
||||
|
||||
@override(Checkpointable)
|
||||
def set_state(self, state: StateDict) -> None:
|
||||
"""Sets the state of the `DifferentiableLearner` instance."""
|
||||
self._check_is_built()
|
||||
|
||||
weights_seq_no = state.get(WEIGHTS_SEQ_NO, 0)
|
||||
|
||||
# Update our weights_seq_no, if the new one is > 0.
|
||||
if weights_seq_no > 0:
|
||||
self._weights_seq_no = weights_seq_no
|
||||
|
||||
# Update our trainable Modules information/function via our config.
|
||||
# If not provided in state (None), all Modules will be trained by default.
|
||||
if "should_module_be_updated" in state:
|
||||
self.config.multi_agent(policies_to_train=state["should_module_be_updated"])
|
||||
|
||||
# TODO (sven): Make `MetricsLogger` a Checkpointable.
|
||||
if COMPONENT_METRICS_LOGGER in state:
|
||||
self.metrics.set_state(state[COMPONENT_METRICS_LOGGER])
|
||||
|
||||
@override(Checkpointable)
|
||||
def get_ctor_args_and_kwargs(self):
|
||||
return (
|
||||
(), # *args,
|
||||
{
|
||||
"config": self.config,
|
||||
}, # **kwargs
|
||||
)
|
||||
|
||||
@override(Checkpointable)
|
||||
def get_checkpointable_components(self):
|
||||
if not self._check_is_built(error=False):
|
||||
self.build()
|
||||
return []
|
||||
|
||||
# TODO (simon): Duplicate in Learner. Move to base class "Learnable".
|
||||
def _make_batch_if_necessary(self, training_data):
|
||||
batch = training_data.batch
|
||||
|
||||
# Call the learner connector on the given `episodes` (if we have one).
|
||||
if training_data.episodes is not None:
|
||||
# If we want to learn from Episodes, we must have a LearnerConnector
|
||||
# pipeline to translate into a train batch first.
|
||||
if self._learner_connector is None:
|
||||
raise ValueError(
|
||||
f"If episodes provided for training, Learner ({self}) must have a "
|
||||
"LearnerConnector pipeline (but pipeline is None)!"
|
||||
)
|
||||
|
||||
# Call the learner connector pipeline.
|
||||
shared_data = {}
|
||||
batch = self._learner_connector(
|
||||
rl_module=self.module,
|
||||
batch=training_data.batch if training_data.batch is not None else {},
|
||||
episodes=training_data.episodes,
|
||||
shared_data=shared_data,
|
||||
metrics=self.metrics,
|
||||
)
|
||||
# Convert to a batch.
|
||||
# TODO (sven): Try to not require MultiAgentBatch anymore.
|
||||
batch = MultiAgentBatch(
|
||||
{
|
||||
module_id: (
|
||||
SampleBatch(module_data, _zero_padded=True)
|
||||
if shared_data.get(f"_zero_padded_for_mid={module_id}")
|
||||
else SampleBatch(module_data)
|
||||
)
|
||||
for module_id, module_data in batch.items()
|
||||
},
|
||||
env_steps=sum(len(e) for e in training_data.episodes),
|
||||
)
|
||||
# Single-agent SampleBatch: Have to convert to MultiAgentBatch.
|
||||
elif isinstance(training_data.batch, SampleBatch):
|
||||
if len(self.module) != 1:
|
||||
raise ValueError(
|
||||
f"SampleBatch provided, but RLModule ({self.module}) has more than "
|
||||
f"one sub-RLModule! Need to provide MultiAgentBatch instead."
|
||||
)
|
||||
|
||||
batch = MultiAgentBatch(
|
||||
{next(iter(self.module.keys())): training_data.batch},
|
||||
env_steps=len(training_data.batch),
|
||||
)
|
||||
# If we already have an `MultiAgentBatch` but with `numpy` array, convert to
|
||||
# tensors.
|
||||
elif (
|
||||
isinstance(training_data.batch, MultiAgentBatch)
|
||||
and training_data.batch.policy_batches
|
||||
and (
|
||||
isinstance(
|
||||
next(iter(training_data.batch.policy_batches.values()))["obs"],
|
||||
numpy.ndarray,
|
||||
)
|
||||
or next(iter(training_data.batch.policy_batches.values()))["obs"].device
|
||||
!= self._device
|
||||
)
|
||||
):
|
||||
batch = self._convert_batch_type(training_data.batch)
|
||||
|
||||
return batch
|
||||
|
||||
# TODO (simon): Duplicate in Learner. Move to base class "Learnable".
|
||||
def _set_slicing_by_batch_id(
|
||||
self, batch: MultiAgentBatch, *, value: bool
|
||||
) -> MultiAgentBatch:
|
||||
"""Enables slicing by batch id in the given batch.
|
||||
|
||||
If the input batch contains batches of sequences we need to make sure when
|
||||
slicing happens it is sliced via batch id and not timestamp. Calling this
|
||||
method enables the same flag on each SampleBatch within the input
|
||||
MultiAgentBatch.
|
||||
|
||||
Args:
|
||||
batch: The MultiAgentBatch to enable slicing by batch id on.
|
||||
value: The value to set the flag to.
|
||||
|
||||
Returns:
|
||||
The input MultiAgentBatch with the indexing flag is enabled / disabled on.
|
||||
"""
|
||||
|
||||
for pid, policy_batch in batch.policy_batches.items():
|
||||
# We assume that arriving batches for recurrent modules OR batches that
|
||||
# have a SEQ_LENS column are already zero-padded to the max sequence length
|
||||
# and have tensors of shape [B, T, ...]. Therefore, we slice sequence
|
||||
# lengths in B. See SampleBatch for more information.
|
||||
if (
|
||||
self.module[pid].is_stateful()
|
||||
or policy_batch.get("seq_lens") is not None
|
||||
):
|
||||
if value:
|
||||
policy_batch.enable_slicing_by_batch_id()
|
||||
else:
|
||||
policy_batch.disable_slicing_by_batch_id()
|
||||
|
||||
return batch
|
||||
|
||||
# TODO (simon): Duplicate in Learner. Move to base class "Learnable".
|
||||
def _check_is_built(self, error: bool = True) -> bool:
|
||||
if self.module is None:
|
||||
if error:
|
||||
raise ValueError(
|
||||
"Learner.build() must be called after constructing a "
|
||||
"Learner and before calling any methods on it."
|
||||
)
|
||||
return False
|
||||
return True
|
||||
|
||||
@abc.abstractmethod
|
||||
def _get_tensor_variable(
|
||||
self,
|
||||
value: Any,
|
||||
dtype: Any = None,
|
||||
trainable: bool = False,
|
||||
) -> TensorType:
|
||||
"""Returns a framework-specific tensor variable with the initial given value.
|
||||
|
||||
This is a framework specific method that should be implemented by the
|
||||
framework specific sub-classes.
|
||||
|
||||
Args:
|
||||
value: The initial value for the tensor variable variable.
|
||||
|
||||
Returns:
|
||||
The framework specific tensor variable of the given initial value,
|
||||
dtype and trainable/requires_grad property.
|
||||
"""
|
||||
|
||||
# TODO (simon): Duplicate in Learner. Move to base class "Learnable".
|
||||
def _reset(self):
|
||||
self.metrics = MetricsLogger(
|
||||
stats_cls_lookup=self.config.stats_cls_lookup,
|
||||
root=False,
|
||||
)
|
||||
self._is_built = False
|
||||
|
||||
# TODO (simon): Duplicate in Learner. Move to base class "Learnable".
|
||||
def _log_steps_trained_metrics(self, batch: MultiAgentBatch):
|
||||
"""Logs this iteration's steps trained, based on given `batch`."""
|
||||
# Loop through all modules.
|
||||
for mid, module_batch in batch.policy_batches.items():
|
||||
# Log weights seq no for this batch.
|
||||
self.metrics.log_value(
|
||||
(mid, WEIGHTS_SEQ_NO),
|
||||
self._weights_seq_no,
|
||||
window=1,
|
||||
)
|
||||
|
||||
module_batch_size = len(module_batch)
|
||||
# Log average batch size (for each module).
|
||||
self.metrics.log_value(
|
||||
key=(mid, MODULE_TRAIN_BATCH_SIZE_MEAN),
|
||||
value=module_batch_size,
|
||||
)
|
||||
# Log module steps (for each module).
|
||||
self.metrics.log_value(
|
||||
key=(mid, NUM_MODULE_STEPS_TRAINED),
|
||||
value=module_batch_size,
|
||||
reduce="sum",
|
||||
)
|
||||
self.metrics.log_value(
|
||||
key=(mid, NUM_MODULE_STEPS_TRAINED_LIFETIME),
|
||||
value=module_batch_size,
|
||||
reduce="lifetime_sum",
|
||||
)
|
||||
# Log module steps (sum of all modules).
|
||||
self.metrics.log_value(
|
||||
key=(ALL_MODULES, NUM_MODULE_STEPS_TRAINED),
|
||||
value=module_batch_size,
|
||||
reduce="sum",
|
||||
)
|
||||
self.metrics.log_value(
|
||||
key=(ALL_MODULES, NUM_MODULE_STEPS_TRAINED_LIFETIME),
|
||||
value=module_batch_size,
|
||||
reduce="lifetime_sum",
|
||||
)
|
||||
# Log env steps (all modules).
|
||||
self.metrics.log_value(
|
||||
(ALL_MODULES, NUM_ENV_STEPS_TRAINED),
|
||||
batch.env_steps(),
|
||||
reduce="sum",
|
||||
)
|
||||
self.metrics.log_value(
|
||||
(ALL_MODULES, NUM_ENV_STEPS_TRAINED_LIFETIME),
|
||||
batch.env_steps(),
|
||||
reduce="lifetime_sum",
|
||||
with_throughput=True,
|
||||
)
|
||||
|
||||
@OverrideToImplementCustomLogic
|
||||
@abc.abstractmethod
|
||||
def _make_functional_call(
|
||||
self,
|
||||
module: MultiRLModule,
|
||||
params: Dict[ModuleID, NamedParamDict],
|
||||
batch: MultiAgentBatch,
|
||||
**kwargs,
|
||||
) -> Dict[ModuleID, Dict[str, TensorType]]:
|
||||
"""Makes a functional call to a module.
|
||||
|
||||
Functional calls enable the Learner to (a) use the same module as its
|
||||
`MetaLearner` and (b) to generate and apply gradients without modifying
|
||||
the `RLModule` parameters directly.
|
||||
|
||||
Args:
|
||||
module: The `MultiRLModule` to be used for the functional call. Note, this
|
||||
module's `forward` method must call the `foward_train`.
|
||||
params: A dictionary containing containing for each `RLModule`'s id its
|
||||
named parameter dictionary. For functional calls to work, these
|
||||
parameters need to be cloned.
|
||||
batch: A `MultiAgentBatch` instance to be used in the functional call.
|
||||
|
||||
Returns:
|
||||
A dictionary with the output of the module's forward pass.
|
||||
"""
|
||||
|
||||
@property
|
||||
def module(self) -> MultiRLModule:
|
||||
"""The MultiRLModule reference that is being used in updates."""
|
||||
return self._module
|
||||
|
||||
@module.setter
|
||||
def module(self, module: MultiRLModule) -> None:
|
||||
"""Sets the `MultiRLModule`.
|
||||
|
||||
Args:
|
||||
module: The reference to the `MultiRLModule` of the class that holds the
|
||||
instance of this `DifferentiableLearner` instance.
|
||||
"""
|
||||
self._module = module
|
||||
@@ -0,0 +1,148 @@
|
||||
from dataclasses import dataclass, fields
|
||||
from typing import Callable, List, Optional, Union
|
||||
|
||||
import gymnasium as gym
|
||||
|
||||
from ray.rllib.connectors.connector_v2 import ConnectorV2
|
||||
from ray.rllib.core.learner.differentiable_learner import DifferentiableLearner
|
||||
from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec
|
||||
from ray.rllib.core.rl_module.rl_module import RLModule
|
||||
from ray.rllib.utils.typing import DeviceType, ModuleID
|
||||
|
||||
|
||||
@dataclass
|
||||
class DifferentiableLearnerConfig:
|
||||
"""Configures a `DifferentiableLearner`."""
|
||||
|
||||
# TODO (simon): We implement only for `PyTorch`, so maybe we use here directly
|
||||
# TorchDifferentiableLearner` and check for this?
|
||||
# The `DifferentiableLearner` class. Must be derived from `DifferentiableLearner`.
|
||||
learner_class: Callable
|
||||
|
||||
learner_connector: Optional[
|
||||
Callable[["RLModule"], Union["ConnectorV2", List["ConnectorV2"]]]
|
||||
] = None
|
||||
|
||||
add_default_connectors_to_learner_pipeline: bool = True
|
||||
|
||||
is_multi_agent: bool = False
|
||||
|
||||
policies_to_update: List[ModuleID] = None
|
||||
|
||||
# The learning rate to use for the nested update. Note, in the default case this
|
||||
# learning rate is only used to update parameters in a functional form, i.e. the
|
||||
# `RLModule`'s stateful parameters are only updated in the `MetaLearner`. Different
|
||||
# logic can be implemented in customized `DifferentiableLearner`s.
|
||||
lr: float = 3e-5
|
||||
|
||||
# TODO (simon): Add further hps like clip_grad, ...
|
||||
# The total number of minibatches to be formed from the batch per learner, e.g.
|
||||
# setting `train_batch_size_per_learner=10` and `num_total_minibatches` to 2
|
||||
# runs 2 SGD minibatch updates with a batch of 5 per training iteration.
|
||||
num_total_minibatches: int = 0
|
||||
|
||||
# The number of epochs per training iteration.
|
||||
num_epochs: int = 1
|
||||
|
||||
# The minibatch size per SGD minibatch update, e.g. with a `train_batch_size_per_learner=10`
|
||||
# and a `minibatch_size=2` the training step runs 5 SGD minibatch updates with minibatches
|
||||
# of 2.
|
||||
minibatch_size: int = None
|
||||
|
||||
# If the batch should be shuffled between epochs.
|
||||
shuffle_batch_per_epoch: bool = False
|
||||
|
||||
def __post_init__(self):
|
||||
"""Additional initialization processes."""
|
||||
|
||||
# Ensure we have a `DifferentiableLearner` class.
|
||||
if not issubclass(self.learner_class, DifferentiableLearner):
|
||||
raise ValueError(
|
||||
"`learner_class` must be a subclass of `DifferentiableLearner "
|
||||
f"but is {self.learner_class}."
|
||||
)
|
||||
|
||||
def build_learner_connector(
|
||||
self,
|
||||
input_observation_space: Optional[gym.spaces.Space],
|
||||
input_action_space: Optional[gym.spaces.Space],
|
||||
device: Optional[DeviceType] = None,
|
||||
):
|
||||
from ray.rllib.connectors.learner import (
|
||||
AddColumnsFromEpisodesToTrainBatch,
|
||||
AddObservationsFromEpisodesToBatch,
|
||||
AddStatesFromEpisodesToBatch,
|
||||
AddTimeDimToBatchAndZeroPad,
|
||||
AgentToModuleMapping,
|
||||
BatchIndividualItems,
|
||||
LearnerConnectorPipeline,
|
||||
NumpyToTensor,
|
||||
)
|
||||
|
||||
custom_connectors = []
|
||||
# Create a learner connector pipeline (including RLlib's default
|
||||
# learner connector piece) and return it.
|
||||
if self.learner_connector is not None:
|
||||
val_ = self.learner_connector(
|
||||
input_observation_space,
|
||||
input_action_space,
|
||||
# device, # TODO (sven): Also pass device into custom builder.
|
||||
)
|
||||
|
||||
from ray.rllib.connectors.connector_v2 import ConnectorV2
|
||||
|
||||
# ConnectorV2 (piece or pipeline).
|
||||
if isinstance(val_, ConnectorV2):
|
||||
custom_connectors = [val_]
|
||||
# Sequence of individual ConnectorV2 pieces.
|
||||
elif isinstance(val_, (list, tuple)):
|
||||
custom_connectors = list(val_)
|
||||
# Unsupported return value.
|
||||
else:
|
||||
raise ValueError(
|
||||
"`AlgorithmConfig.training(learner_connector=..)` must return "
|
||||
"a ConnectorV2 object or a list thereof (to be added to a "
|
||||
f"pipeline)! Your function returned {val_}."
|
||||
)
|
||||
|
||||
pipeline = LearnerConnectorPipeline(
|
||||
connectors=custom_connectors,
|
||||
input_observation_space=input_observation_space,
|
||||
input_action_space=input_action_space,
|
||||
)
|
||||
if self.add_default_connectors_to_learner_pipeline:
|
||||
# Append OBS handling.
|
||||
pipeline.append(
|
||||
AddObservationsFromEpisodesToBatch(as_learner_connector=True)
|
||||
)
|
||||
# Append all other columns handling.
|
||||
pipeline.append(AddColumnsFromEpisodesToTrainBatch())
|
||||
# Append time-rank handler.
|
||||
pipeline.append(AddTimeDimToBatchAndZeroPad(as_learner_connector=True))
|
||||
# Append STATE_IN/STATE_OUT handler.
|
||||
pipeline.append(AddStatesFromEpisodesToBatch(as_learner_connector=True))
|
||||
# If multi-agent -> Map from AgentID-based data to ModuleID based data.
|
||||
if self.is_multi_agent:
|
||||
pipeline.append(
|
||||
AgentToModuleMapping(
|
||||
rl_module_specs=(
|
||||
self.rl_module_spec.rl_module_specs
|
||||
if isinstance(self.rl_module_spec, MultiRLModuleSpec)
|
||||
else set(self.policies)
|
||||
),
|
||||
agent_to_module_mapping_fn=self.policy_mapping_fn,
|
||||
)
|
||||
)
|
||||
# Batch all data.
|
||||
pipeline.append(BatchIndividualItems(multi_agent=self.is_multi_agent))
|
||||
# Convert to Tensors.
|
||||
pipeline.append(NumpyToTensor(as_learner_connector=True, device=device))
|
||||
return pipeline
|
||||
|
||||
def update_from_kwargs(self, **kwargs):
|
||||
"""Sets all slots with values defined in `kwargs`."""
|
||||
# Get all field names (i.e., slot names).
|
||||
field_names = {f.name for f in fields(self)}
|
||||
for key, value in kwargs.items():
|
||||
if key in field_names:
|
||||
setattr(self, key, value)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,821 @@
|
||||
import copy
|
||||
import itertools
|
||||
from functools import partial
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Callable,
|
||||
Collection,
|
||||
Dict,
|
||||
List,
|
||||
Optional,
|
||||
Set,
|
||||
Type,
|
||||
Union,
|
||||
)
|
||||
|
||||
import ray
|
||||
from ray._common.deprecation import Deprecated
|
||||
from ray.rllib.core import (
|
||||
COMPONENT_LEARNER,
|
||||
COMPONENT_RL_MODULE,
|
||||
)
|
||||
from ray.rllib.core.learner.learner import Learner
|
||||
from ray.rllib.core.learner.training_data import TrainingData
|
||||
from ray.rllib.core.rl_module import validate_module_id
|
||||
from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec
|
||||
from ray.rllib.core.rl_module.rl_module import RLModuleSpec
|
||||
from ray.rllib.policy.policy import PolicySpec
|
||||
from ray.rllib.policy.sample_batch import MultiAgentBatch
|
||||
from ray.rllib.utils.actor_manager import (
|
||||
FaultTolerantActorManager,
|
||||
RemoteCallResults,
|
||||
ResultOrError,
|
||||
)
|
||||
from ray.rllib.utils.annotations import override
|
||||
from ray.rllib.utils.checkpoints import Checkpointable
|
||||
from ray.rllib.utils.metrics.ray_metrics import (
|
||||
DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
|
||||
TimerAndPrometheusLogger,
|
||||
)
|
||||
from ray.rllib.utils.typing import (
|
||||
EpisodeType,
|
||||
ModuleID,
|
||||
RLModuleSpecType,
|
||||
ShouldModuleBeUpdatedFn,
|
||||
StateDict,
|
||||
T,
|
||||
)
|
||||
from ray.train._internal.backend_executor import BackendExecutor
|
||||
from ray.util.annotations import PublicAPI
|
||||
from ray.util.metrics import Histogram
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
|
||||
from ray.util.placement_group import PlacementGroup
|
||||
|
||||
|
||||
def _get_backend_config(learner_class: Type[Learner]) -> str:
|
||||
if learner_class.framework == "torch":
|
||||
from ray.train.torch.config import TorchConfig, _TorchBackend
|
||||
|
||||
# Override `_TorchBackend` share_cuda_visible_devices=True setting.
|
||||
# We need this to be False to make sure Learner actors only see their
|
||||
# own GPU. There is no need in RLlib's LearnerGroups for 2 different Learner
|
||||
# actors to communicate with each other through their GPUs.
|
||||
class _RLlibTorchBackend(_TorchBackend):
|
||||
share_cuda_visible_devices = False
|
||||
|
||||
class RLlibTorchConfig(TorchConfig):
|
||||
@property
|
||||
def backend_cls(self):
|
||||
return _RLlibTorchBackend
|
||||
|
||||
backend_config = RLlibTorchConfig()
|
||||
|
||||
else:
|
||||
raise ValueError(
|
||||
"`learner_class.framework` must be 'torch' (but is "
|
||||
f"{learner_class.framework}!"
|
||||
)
|
||||
|
||||
return backend_config
|
||||
|
||||
|
||||
class RLlibBackendExecutor(BackendExecutor):
|
||||
# Override `BackendExecutor` placement group creation logic. We need to pass our own
|
||||
# to make sure the one of the Algorithm (Trainable) is used for all the
|
||||
# Algorithm's actors.
|
||||
def _create_placement_group(self):
|
||||
pass
|
||||
|
||||
# TODO (sven): Change this once there is a better (public) API for this in the
|
||||
# superclass.
|
||||
def set_placement_group(self, placement_group):
|
||||
if placement_group is not None:
|
||||
self._placement_group = placement_group
|
||||
|
||||
|
||||
@PublicAPI(stability="alpha")
|
||||
class LearnerGroup(Checkpointable):
|
||||
"""Coordinator of n (possibly remote) Learner workers.
|
||||
|
||||
Each Learner worker has a copy of the RLModule, the loss function(s), and
|
||||
one or more optimizers.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
config: "AlgorithmConfig",
|
||||
# TODO (sven): Rename into `rl_module_spec`.
|
||||
module_spec: Optional[RLModuleSpecType] = None,
|
||||
placement_group: Optional["PlacementGroup"] = None,
|
||||
):
|
||||
"""Initializes a LearnerGroup instance.
|
||||
|
||||
Args:
|
||||
config: The AlgorithmConfig object to use to configure this LearnerGroup.
|
||||
Call the `learners(num_learners=...)` method on your config to
|
||||
specify the number of learner workers to use.
|
||||
Call the same method with arguments `num_cpus_per_learner` and/or
|
||||
`num_gpus_per_learner` to configure the compute used by each
|
||||
Learner worker in this LearnerGroup.
|
||||
Call the `training(learner_class=...)` method on your config to specify,
|
||||
which exact Learner class to use.
|
||||
Call the `rl_module(rl_module_spec=...)` method on your config to set up
|
||||
the specifics for your RLModule to be used in each Learner.
|
||||
module_spec: If not already specified in `config`, a separate overriding
|
||||
RLModuleSpec may be provided via this argument.
|
||||
placement_group: An optional `PlacementGroup` instance to set the
|
||||
`RLlibBackendExecutor`'s `self._placement_group` attribute to.
|
||||
If run within an Algorithm (tune.Trainable), the placement group of tune
|
||||
trial actor is passed through here.
|
||||
"""
|
||||
self.config = config.copy(copy_frozen=False)
|
||||
self._module_spec = module_spec
|
||||
|
||||
learner_class = self.config.learner_class
|
||||
module_spec = module_spec or self.config.get_multi_rl_module_spec()
|
||||
|
||||
self._learner = None
|
||||
self._workers = None
|
||||
# If a user calls self.shutdown() on their own then this flag is set to true.
|
||||
# When del is called the backend executor isn't shutdown twice if this flag is
|
||||
# true. the backend executor would otherwise log a warning to the console from
|
||||
# ray train.
|
||||
self._is_shut_down = False
|
||||
|
||||
# How many timesteps had to be dropped due to a full input queue?
|
||||
self._ts_dropped = 0
|
||||
|
||||
# A single local Learner.
|
||||
if not self.is_remote:
|
||||
self._learner = learner_class(config=config, module_spec=module_spec)
|
||||
self._learner.build()
|
||||
self._worker_manager = None
|
||||
|
||||
# Ray metrics
|
||||
self._metrics_local_learner_training_data_solve_refs = Histogram(
|
||||
name="rllib_learner_local_training_data_solve_refs_time",
|
||||
description="Time spent in resolve training data refs for local learner.",
|
||||
boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
|
||||
tag_keys=("rllib",),
|
||||
)
|
||||
self._metrics_local_learner_training_data_solve_refs.set_default_tags(
|
||||
{"rllib": self.__class__.__name__}
|
||||
)
|
||||
|
||||
# N remote Learner workers.
|
||||
else:
|
||||
backend_config = _get_backend_config(learner_class)
|
||||
|
||||
num_cpus_per_learner = (
|
||||
self.config.num_cpus_per_learner
|
||||
if self.config.num_cpus_per_learner != "auto"
|
||||
else 1
|
||||
if self.config.num_gpus_per_learner == 0
|
||||
else 0
|
||||
)
|
||||
num_gpus_per_learner = max(0, self.config.num_gpus_per_learner)
|
||||
resources_per_learner = {
|
||||
"CPU": num_cpus_per_learner,
|
||||
"GPU": num_gpus_per_learner,
|
||||
**(self.config.custom_resources_per_learner or {}),
|
||||
}
|
||||
|
||||
backend_executor = RLlibBackendExecutor(
|
||||
backend_config=backend_config,
|
||||
num_workers=self.config.num_learners,
|
||||
resources_per_worker=resources_per_learner,
|
||||
max_retries=0,
|
||||
)
|
||||
# Set the placement group - if any - of the BackendExecutor.
|
||||
backend_executor.set_placement_group(placement_group)
|
||||
backend_executor.start(
|
||||
train_cls=learner_class,
|
||||
train_cls_kwargs={
|
||||
"config": config,
|
||||
"module_spec": module_spec,
|
||||
},
|
||||
)
|
||||
self._backend_executor = backend_executor
|
||||
|
||||
self._workers = [w.actor for w in backend_executor.worker_group.workers]
|
||||
|
||||
ray.get(
|
||||
[
|
||||
worker._set_learner_index_and_placement_group.remote(
|
||||
learner_index=idx,
|
||||
placement_group=placement_group,
|
||||
)
|
||||
for idx, worker in enumerate(self._workers)
|
||||
]
|
||||
)
|
||||
|
||||
# Run the neural network building code on remote workers.
|
||||
ray.get([w.build.remote() for w in self._workers])
|
||||
|
||||
self._worker_manager = FaultTolerantActorManager(
|
||||
self._workers,
|
||||
max_remote_requests_in_flight_per_actor=(
|
||||
self.config.max_requests_in_flight_per_learner
|
||||
),
|
||||
)
|
||||
|
||||
# Ray metrics
|
||||
self._metrics_update_time = Histogram(
|
||||
name="rllib_learner_group_update_time",
|
||||
description="Time spent in LearnerGroup.update()",
|
||||
boundaries=DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS,
|
||||
tag_keys=("rllib",),
|
||||
)
|
||||
self._metrics_update_time.set_default_tags({"rllib": self.__class__.__name__})
|
||||
|
||||
# TODO (sven): Replace this with call to `self.metrics.peek()`?
|
||||
# Currently LearnerGroup does not have a metrics object.
|
||||
def get_stats(self) -> Dict[str, Any]:
|
||||
"""Returns the current stats for the input queue for this learner group."""
|
||||
return {
|
||||
"learner_group_ts_dropped_lifetime": self._ts_dropped,
|
||||
"actor_manager_num_outstanding_async_reqs": (
|
||||
0
|
||||
if self.is_local
|
||||
else self._worker_manager.num_outstanding_async_reqs()
|
||||
),
|
||||
}
|
||||
|
||||
@property
|
||||
def is_remote(self) -> bool:
|
||||
return self.config.num_learners > 0
|
||||
|
||||
@property
|
||||
def is_local(self) -> bool:
|
||||
return not self.is_remote
|
||||
|
||||
def update(
|
||||
self,
|
||||
*,
|
||||
batch: Optional[MultiAgentBatch] = None,
|
||||
batches: Optional[List[MultiAgentBatch]] = None,
|
||||
batch_refs: Optional[List[ray.ObjectRef]] = None,
|
||||
episodes: Optional[List[EpisodeType]] = None,
|
||||
episodes_refs: Optional[List[ray.ObjectRef]] = None,
|
||||
data_iterators: Optional[List[ray.data.DataIterator]] = None,
|
||||
training_data: Optional[TrainingData] = None,
|
||||
timesteps: Optional[Dict[str, Any]] = None,
|
||||
async_update: bool = False,
|
||||
return_state: bool = False,
|
||||
# User kwargs passed onto the Learners.
|
||||
**kwargs,
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Performs gradient based updates on Learners in parallel.
|
||||
|
||||
Updates are performed with data from any of the provided arguments
|
||||
(batch, batches, batch_refs, episodes, episodes_refs, data_iterators, training_data).
|
||||
|
||||
Args:
|
||||
batch: A data batch to use for the update. If there are more
|
||||
than one Learner workers, the batch is split amongst these and one
|
||||
shard is sent to each Learner.
|
||||
batch_refs: A list of Ray ObjectRefs to the batches. If there are more
|
||||
than one Learner workers, the list of batch refs is split amongst these and
|
||||
one list shard is sent to each Learner.
|
||||
episodes: A list of Episodes to process and perform the update
|
||||
for. If there are more than one Learner workers, the list of episodes
|
||||
is split amongst these and one list shard is sent to each Learner.
|
||||
episodes_refs: A list of Ray ObjectRefs to the episodes. If there are more
|
||||
than one Learner workers, the list of episode refs is split amongst these and
|
||||
one list shard is sent to each Learner.
|
||||
timesteps: A dictionary of timesteps to pass to the Learners's update method.
|
||||
This is usually used for learning rate scheduling but can be used for any other purpose.
|
||||
training_data: A TrainingData object to use for the update. If not provided,
|
||||
a new TrainingData object will be created from the batch, batches, batch_refs,
|
||||
episodes, and episodes_refs.
|
||||
async_update: Whether the update request(s) to the Learner workers should be
|
||||
sent asynchronously. If True, will return NOT the results from the
|
||||
update on the given data, but all results from prior asynchronous update
|
||||
requests that have not been returned thus far.
|
||||
return_state: Whether to include one of the Learner worker's state from
|
||||
after the update step in the returned results dict (under the
|
||||
`_rl_module_state_after_update` key). Note that after an update, all
|
||||
Learner workers' states should be identical, so we use the first
|
||||
Learner's state here. Useful for avoiding an extra `get_weights()` call,
|
||||
e.g. for synchronizing EnvRunner weights.
|
||||
num_epochs: The number of complete passes over the entire train batch. Each
|
||||
pass might be further split into n minibatches (if `minibatch_size`
|
||||
provided).
|
||||
minibatch_size: The size of minibatches to use to further split the train
|
||||
`batch` into sub-batches. The `batch` is then iterated over n times
|
||||
where n is `len(batch) // minibatch_size`.
|
||||
shuffle_batch_per_epoch: Whether to shuffle the train batch once per epoch.
|
||||
If the train batch has a time rank (axis=1), shuffling will only take
|
||||
place along the batch axis to not disturb any intact (episode)
|
||||
trajectories. Also, shuffling is always skipped if `minibatch_size` is
|
||||
None, meaning the entire train batch is processed each epoch, making it
|
||||
unnecessary to shuffle.
|
||||
**kwargs:
|
||||
|
||||
Returns:
|
||||
If `async_update` is False, a dictionary with the reduced results of the
|
||||
updates from the Learner(s) or a list of dictionaries of results from the
|
||||
updates from the Learner(s).
|
||||
If `async_update` is True, a list of list of dictionaries of results, where
|
||||
the outer list corresponds to separate previous calls to this method, and
|
||||
the inner list corresponds to the results from each Learner(s). Or if the
|
||||
results are reduced, a list of dictionaries of the reduced results from each
|
||||
call to async_update that is ready.
|
||||
"""
|
||||
with TimerAndPrometheusLogger(self._metrics_update_time):
|
||||
# Create and validate TrainingData object, if not already provided.
|
||||
if training_data is None:
|
||||
training_data = TrainingData(
|
||||
batch=batch,
|
||||
batches=batches,
|
||||
batch_refs=batch_refs,
|
||||
episodes=episodes,
|
||||
episodes_refs=episodes_refs,
|
||||
data_iterators=data_iterators,
|
||||
)
|
||||
training_data.validate()
|
||||
|
||||
# NEW: allow caller to defer Ray.get()/materialization to the learner thread.
|
||||
# TODO (simon): Set to `False` and create attribute in config.
|
||||
defer_solve = kwargs.pop("defer_solve_refs_to_learner", False)
|
||||
|
||||
# Local Learner instance.
|
||||
if self.is_local:
|
||||
if async_update:
|
||||
raise ValueError(
|
||||
"Can't call `update(async_update=True)` when running with "
|
||||
"`num_learners=0`! Set `config.num_learners > 0` to allow async "
|
||||
"updates."
|
||||
)
|
||||
# Only solve refs here if NOT deferring. When deferring, the Learner/GPU
|
||||
# loader thread will call `training_data.solve_refs()` and build the CPU MAB.
|
||||
if not defer_solve:
|
||||
# Ray metrics
|
||||
with TimerAndPrometheusLogger(
|
||||
self._metrics_local_learner_training_data_solve_refs
|
||||
):
|
||||
training_data.solve_refs()
|
||||
|
||||
if return_state:
|
||||
kwargs["return_state"] = return_state
|
||||
# Return the single Learner's update results.
|
||||
return [
|
||||
self._learner.update(
|
||||
training_data=training_data,
|
||||
timesteps=timesteps,
|
||||
**kwargs,
|
||||
)
|
||||
]
|
||||
|
||||
# Remote Learner actors' kwargs.
|
||||
remote_call_kwargs = [
|
||||
dict(
|
||||
training_data=td_shard,
|
||||
timesteps=timesteps,
|
||||
# If `return_state=True`, only return it from the first Learner
|
||||
# actor.
|
||||
return_state=(return_state and i == 0),
|
||||
**kw,
|
||||
**kwargs,
|
||||
)
|
||||
for i, (td_shard, kw) in enumerate(
|
||||
training_data.shard(
|
||||
num_shards=len(self),
|
||||
len_lookback_buffer=self.config.episode_lookback_horizon,
|
||||
**kwargs,
|
||||
)
|
||||
)
|
||||
]
|
||||
|
||||
# Async updates.
|
||||
if async_update:
|
||||
# Retrieve all ready results (kicked off by prior calls to this method).
|
||||
results = self._worker_manager.fetch_ready_async_reqs(
|
||||
timeout_seconds=0.0
|
||||
)
|
||||
# Send out new request(s), if there is still capacity on the actors
|
||||
# (each actor is allowed only some number of max in-flight requests
|
||||
# at the same time).
|
||||
num_sent_requests = self._worker_manager.foreach_actor_async(
|
||||
"update",
|
||||
kwargs=remote_call_kwargs,
|
||||
)
|
||||
|
||||
# Some requests were dropped, record lost ts/data.
|
||||
if num_sent_requests != len(self):
|
||||
factor = 1 - (num_sent_requests / len(self))
|
||||
# TODO (sven): Move this logic into a TrainingData API as well
|
||||
# (`TrainingData.env_steps()`).
|
||||
if training_data.batch_refs is not None:
|
||||
dropped = (
|
||||
len(training_data.batch_refs)
|
||||
* self.config.train_batch_size_per_learner
|
||||
)
|
||||
elif training_data.batch is not None:
|
||||
dropped = len(training_data.batch)
|
||||
# List of Ray ObjectRefs (each object ref is a list of episodes of
|
||||
# total len=`rollout_fragment_length * num_envs_per_env_runner`)
|
||||
elif training_data.episodes_refs is not None:
|
||||
dropped = (
|
||||
len(training_data.episodes_refs)
|
||||
* self.config.get_rollout_fragment_length()
|
||||
* self.config.num_envs_per_env_runner
|
||||
)
|
||||
else:
|
||||
assert training_data.episodes is not None
|
||||
dropped = sum(len(e) for e in training_data.episodes)
|
||||
|
||||
self._ts_dropped += factor * dropped
|
||||
# Sync updates.
|
||||
else:
|
||||
results = self._worker_manager.foreach_actor(
|
||||
"update",
|
||||
kwargs=remote_call_kwargs,
|
||||
)
|
||||
|
||||
results = self._get_results(results)
|
||||
|
||||
return results
|
||||
|
||||
def add_module(
|
||||
self,
|
||||
*,
|
||||
module_id: ModuleID,
|
||||
module_spec: RLModuleSpec,
|
||||
config_overrides: Optional[Dict] = None,
|
||||
new_should_module_be_updated: Optional[ShouldModuleBeUpdatedFn] = None,
|
||||
) -> MultiRLModuleSpec:
|
||||
"""Adds a module to the underlying MultiRLModule.
|
||||
|
||||
Changes this Learner's config in order to make this architectural change
|
||||
permanent wrt. to checkpointing.
|
||||
|
||||
Args:
|
||||
module_id: The ModuleID of the module to be added.
|
||||
module_spec: The ModuleSpec of the module to be added.
|
||||
config_overrides: The `AlgorithmConfig` overrides that should apply to
|
||||
the new Module, if any.
|
||||
new_should_module_be_updated: An optional sequence of ModuleIDs or a
|
||||
callable taking ModuleID and SampleBatchType and returning whether the
|
||||
ModuleID should be updated (trained).
|
||||
If None, will keep the existing setup in place. RLModules,
|
||||
whose IDs are not in the list (or for which the callable
|
||||
returns False) will not be updated.
|
||||
|
||||
Returns:
|
||||
The new MultiRLModuleSpec (after the change has been performed).
|
||||
"""
|
||||
validate_module_id(module_id, error=True)
|
||||
|
||||
# Force-set inference-only = False.
|
||||
module_spec = copy.deepcopy(module_spec)
|
||||
module_spec.inference_only = False
|
||||
|
||||
results = self.foreach_learner(
|
||||
func=lambda _learner: _learner.add_module(
|
||||
module_id=module_id,
|
||||
module_spec=module_spec,
|
||||
config_overrides=config_overrides,
|
||||
new_should_module_be_updated=new_should_module_be_updated,
|
||||
),
|
||||
)
|
||||
marl_spec = self._get_results(results)[0]
|
||||
|
||||
# Change our config (AlgorithmConfig) to contain the new Module.
|
||||
# TODO (sven): This is a hack to manipulate the AlgorithmConfig directly,
|
||||
# but we'll deprecate config.policies soon anyway.
|
||||
self.config.policies[module_id] = PolicySpec()
|
||||
if config_overrides is not None:
|
||||
self.config.multi_agent(
|
||||
algorithm_config_overrides_per_module={module_id: config_overrides}
|
||||
)
|
||||
self.config.rl_module(rl_module_spec=marl_spec)
|
||||
if new_should_module_be_updated is not None:
|
||||
self.config.multi_agent(policies_to_train=new_should_module_be_updated)
|
||||
|
||||
return marl_spec
|
||||
|
||||
def remove_module(
|
||||
self,
|
||||
module_id: ModuleID,
|
||||
*,
|
||||
new_should_module_be_updated: Optional[ShouldModuleBeUpdatedFn] = None,
|
||||
) -> MultiRLModuleSpec:
|
||||
"""Removes a module from the Learner.
|
||||
|
||||
Args:
|
||||
module_id: The ModuleID of the module to be removed.
|
||||
new_should_module_be_updated: An optional sequence of ModuleIDs or a
|
||||
callable taking ModuleID and SampleBatchType and returning whether the
|
||||
ModuleID should be updated (trained).
|
||||
If None, will keep the existing setup in place. RLModules,
|
||||
whose IDs are not in the list (or for which the callable
|
||||
returns False) will not be updated.
|
||||
|
||||
Returns:
|
||||
The new MultiRLModuleSpec (after the change has been performed).
|
||||
"""
|
||||
results = self.foreach_learner(
|
||||
func=lambda _learner: _learner.remove_module(
|
||||
module_id=module_id,
|
||||
new_should_module_be_updated=new_should_module_be_updated,
|
||||
),
|
||||
)
|
||||
marl_spec = self._get_results(results)[0]
|
||||
|
||||
# Change self.config to reflect the new architecture.
|
||||
# TODO (sven): This is a hack to manipulate the AlgorithmConfig directly,
|
||||
# but we'll deprecate config.policies soon anyway.
|
||||
del self.config.policies[module_id]
|
||||
self.config.algorithm_config_overrides_per_module.pop(module_id, None)
|
||||
if new_should_module_be_updated is not None:
|
||||
self.config.multi_agent(policies_to_train=new_should_module_be_updated)
|
||||
self.config.rl_module(rl_module_spec=marl_spec)
|
||||
|
||||
return marl_spec
|
||||
|
||||
@override(Checkpointable)
|
||||
def get_state(
|
||||
self,
|
||||
components: Optional[Union[str, Collection[str]]] = None,
|
||||
*,
|
||||
not_components: Optional[Union[str, Collection[str]]] = None,
|
||||
**kwargs,
|
||||
) -> StateDict:
|
||||
state = {}
|
||||
|
||||
if self._check_component(COMPONENT_LEARNER, components, not_components):
|
||||
if self.is_local:
|
||||
state[COMPONENT_LEARNER] = self._learner.get_state(
|
||||
components=self._get_subcomponents(COMPONENT_LEARNER, components),
|
||||
not_components=self._get_subcomponents(
|
||||
COMPONENT_LEARNER, not_components
|
||||
),
|
||||
**kwargs,
|
||||
)
|
||||
else:
|
||||
worker = self._worker_manager.healthy_actor_ids()[0]
|
||||
assert len(self) == self._worker_manager.num_healthy_actors()
|
||||
_comps = self._get_subcomponents(COMPONENT_LEARNER, components)
|
||||
_not_comps = self._get_subcomponents(COMPONENT_LEARNER, not_components)
|
||||
results = self._worker_manager.foreach_actor(
|
||||
lambda w: w.get_state(_comps, not_components=_not_comps, **kwargs),
|
||||
remote_actor_ids=[worker],
|
||||
)
|
||||
state[COMPONENT_LEARNER] = self._get_results(results)[0]
|
||||
|
||||
return state
|
||||
|
||||
@override(Checkpointable)
|
||||
def set_state(self, state: StateDict) -> None:
|
||||
if COMPONENT_LEARNER in state:
|
||||
if self.is_local:
|
||||
self._learner.set_state(state[COMPONENT_LEARNER])
|
||||
else:
|
||||
state_ref = ray.put(state[COMPONENT_LEARNER])
|
||||
self.foreach_learner(
|
||||
lambda _learner, _ref=state_ref: _learner.set_state(ray.get(_ref))
|
||||
)
|
||||
|
||||
def get_weights(
|
||||
self, module_ids: Optional[Collection[ModuleID]] = None
|
||||
) -> StateDict:
|
||||
"""Convenience method instead of self.get_state(components=...).
|
||||
|
||||
Args:
|
||||
module_ids: An optional collection of ModuleIDs for which to return weights.
|
||||
If None (default), return weights of all RLModules.
|
||||
|
||||
Returns:
|
||||
The results of
|
||||
`self.get_state(components='learner/rl_module')['learner']['rl_module']`.
|
||||
"""
|
||||
# Return the entire RLModule state (all possible single-agent RLModules).
|
||||
if module_ids is None:
|
||||
components = COMPONENT_LEARNER + "/" + COMPONENT_RL_MODULE
|
||||
# Return a subset of the single-agent RLModules.
|
||||
else:
|
||||
components = [
|
||||
"".join(tup)
|
||||
for tup in itertools.product(
|
||||
[COMPONENT_LEARNER + "/" + COMPONENT_RL_MODULE + "/"],
|
||||
list(module_ids),
|
||||
)
|
||||
]
|
||||
state = self.get_state(components)[COMPONENT_LEARNER][COMPONENT_RL_MODULE]
|
||||
return state
|
||||
|
||||
def set_weights(self, weights) -> None:
|
||||
"""Convenience method instead of self.set_state({'learner': {'rl_module': ..}}).
|
||||
|
||||
Args:
|
||||
weights: The weights dict of the MultiRLModule of a Learner inside this
|
||||
LearnerGroup.
|
||||
"""
|
||||
self.set_state({COMPONENT_LEARNER: {COMPONENT_RL_MODULE: weights}})
|
||||
|
||||
@override(Checkpointable)
|
||||
def get_ctor_args_and_kwargs(self):
|
||||
return (
|
||||
(), # *args
|
||||
{
|
||||
"config": self.config,
|
||||
"module_spec": self._module_spec,
|
||||
}, # **kwargs
|
||||
)
|
||||
|
||||
@override(Checkpointable)
|
||||
def get_checkpointable_components(self):
|
||||
# Return the entire ActorManager, if remote. Otherwise, return the
|
||||
# local worker. Also, don't give the component (Learner) a name ("")
|
||||
# as it's the only component in this LearnerGroup to be saved.
|
||||
return [
|
||||
(
|
||||
COMPONENT_LEARNER,
|
||||
self._learner if self.is_local else self._worker_manager,
|
||||
)
|
||||
]
|
||||
|
||||
def foreach_learner(
|
||||
self,
|
||||
func: Callable[[Learner, Optional[Any]], T],
|
||||
*,
|
||||
healthy_only: bool = True,
|
||||
remote_actor_ids: List[int] = None,
|
||||
timeout_seconds: Optional[float] = None,
|
||||
return_obj_refs: bool = False,
|
||||
mark_healthy: bool = False,
|
||||
**kwargs,
|
||||
) -> RemoteCallResults:
|
||||
r"""Calls the given function on each Learner L with the args: (L, \*\*kwargs).
|
||||
|
||||
Args:
|
||||
func: The function to call on each Learner L with args: (L, \*\*kwargs).
|
||||
healthy_only: If True, applies `func` only to Learner actors currently
|
||||
tagged "healthy", otherwise to all actors. If `healthy_only=False` and
|
||||
`mark_healthy=True`, will send `func` to all actors and mark those
|
||||
actors "healthy" that respond to the request within `timeout_seconds`
|
||||
and are currently tagged as "unhealthy".
|
||||
remote_actor_ids: Apply func on a selected set of remote actors. Use None
|
||||
(default) for all actors.
|
||||
timeout_seconds: Time to wait (in seconds) for results. Set this to 0.0 for
|
||||
fire-and-forget. Set this to None (default) to wait infinitely (i.e. for
|
||||
synchronous execution).
|
||||
return_obj_refs: whether to return ObjectRef instead of actual results.
|
||||
Note, for fault tolerance reasons, these returned ObjectRefs should
|
||||
never be resolved with ray.get() outside of the context of this manager.
|
||||
mark_healthy: Whether to mark all those actors healthy again that are
|
||||
currently marked unhealthy AND that returned results from the remote
|
||||
call (within the given `timeout_seconds`).
|
||||
Note that actors are NOT set unhealthy, if they simply time out
|
||||
(only if they return a RayActorError).
|
||||
Also not that this setting is ignored if `healthy_only=True` (b/c this
|
||||
setting only affects actors that are currently tagged as unhealthy).
|
||||
|
||||
Returns:
|
||||
A list of size len(Learners) with the return values of all calls to `func`.
|
||||
"""
|
||||
if self.is_local:
|
||||
results = RemoteCallResults()
|
||||
results.add_result(
|
||||
None,
|
||||
ResultOrError(result=func(self._learner, **kwargs)),
|
||||
None,
|
||||
)
|
||||
return results
|
||||
|
||||
return self._worker_manager.foreach_actor(
|
||||
func=partial(func, **kwargs),
|
||||
healthy_only=healthy_only,
|
||||
remote_actor_ids=remote_actor_ids,
|
||||
timeout_seconds=timeout_seconds,
|
||||
return_obj_refs=return_obj_refs,
|
||||
mark_healthy=mark_healthy,
|
||||
)
|
||||
|
||||
def __len__(self):
|
||||
return 0 if self.is_local else len(self._workers)
|
||||
|
||||
def shutdown(self):
|
||||
"""Shuts down the LearnerGroup."""
|
||||
if self.is_local and self._learner is not None:
|
||||
self._learner.shutdown()
|
||||
if self.is_remote and hasattr(self, "_backend_executor"):
|
||||
self._backend_executor.shutdown(graceful_termination=True)
|
||||
self._is_shut_down = True
|
||||
|
||||
def __del__(self):
|
||||
if not self._is_shut_down:
|
||||
self.shutdown()
|
||||
|
||||
def _get_results(self, results):
|
||||
processed_results = []
|
||||
for result in results:
|
||||
result_or_error = result.get()
|
||||
if result.ok:
|
||||
processed_results.append(result_or_error)
|
||||
else:
|
||||
raise result_or_error
|
||||
return processed_results
|
||||
|
||||
@Deprecated(new="LearnerGroup.update(batch=.., **kwargs)", error=False)
|
||||
def update_from_batch(self, batch, **kwargs):
|
||||
return self.update(batch=batch, **kwargs)
|
||||
|
||||
@Deprecated(new="LearnerGroup.update(episodes=.., **kwargs)", error=False)
|
||||
def update_from_episodes(self, episodes, **kwargs):
|
||||
return self.update(episodes=episodes, **kwargs)
|
||||
|
||||
@Deprecated(new="LearnerGroup.update_from_batch(async=True)", error=True)
|
||||
def async_update(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
@Deprecated(
|
||||
old="LearnerGroup.load_module_state()",
|
||||
help="To restore RLModule or MultiRLModule state "
|
||||
"use LearnerGroup.restore_from_path(path=..., component=...). "
|
||||
"See docs for more details: "
|
||||
"https://docs.ray.io/en/latest/rllib/rl-modules.html#checkpointing-rlmodules",
|
||||
error=False,
|
||||
)
|
||||
def load_module_state(
|
||||
self,
|
||||
*,
|
||||
multi_rl_module_ckpt_dir: Optional[str] = None,
|
||||
modules_to_load: Optional[Set[str]] = None,
|
||||
rl_module_ckpt_dirs: Optional[Dict[ModuleID, str]] = None,
|
||||
) -> None:
|
||||
"""Load the checkpoints of the modules being trained by `LearnerGroup`.
|
||||
|
||||
`load_module_state` can be used 3 ways:
|
||||
1. Load a checkpoint for the `MultiRLModule` being trained by this
|
||||
`LearnerGroup`. Optionally, limit the modules that are loaded
|
||||
from the checkpoint by specifying the `modules_to_load` argument.
|
||||
2. Load the checkpoint(s) for single agent `RLModules` that
|
||||
are in the `MultiRLModule` being trained by this `LearnerGroup`.
|
||||
3. Load a checkpoint for the `MultiRLModule` being trained by this
|
||||
`LearnerGroup` and load the checkpoint(s) for single agent `RLModules`
|
||||
that are in the `MultiRLModule`. The checkpoints for the single
|
||||
agent `RLModules` take precedence over the module states in the
|
||||
`MultiRLModule` checkpoint.
|
||||
|
||||
At least one of `multi_rl_module_ckpt_dir` or `rl_module_ckpt_dirs`
|
||||
must be specified.
|
||||
`modules_to_load` can only be specified if `multi_rl_module_ckpt_dir`
|
||||
is provided.
|
||||
|
||||
Args:
|
||||
multi_rl_module_ckpt_dir: The path to the checkpoint for the
|
||||
`MultiRLModule`.
|
||||
modules_to_load: A set of `RLModule` ids to load from the checkpoint.
|
||||
rl_module_ckpt_dirs: A mapping from module ids to the path to a
|
||||
checkpoint for a single agent `RLModule`.
|
||||
"""
|
||||
if not (multi_rl_module_ckpt_dir or rl_module_ckpt_dirs):
|
||||
raise ValueError(
|
||||
f"At least one of `multi_rl_module_ckpt_dir` or "
|
||||
f"`rl_module_ckpt_dirs` must be provided. "
|
||||
f"Got {multi_rl_module_ckpt_dir=} and {rl_module_ckpt_dirs=}."
|
||||
)
|
||||
|
||||
if modules_to_load and not multi_rl_module_ckpt_dir:
|
||||
raise ValueError(
|
||||
f"`modules_to_load` can only be specified if a "
|
||||
f"multi_rl_module_ckpt_dir is provided. "
|
||||
f"Got {modules_to_load=} and {multi_rl_module_ckpt_dir=}."
|
||||
)
|
||||
|
||||
# MultiRLModule checkpoint is provided.
|
||||
if multi_rl_module_ckpt_dir:
|
||||
# Restore the entire MultiRLModule state.
|
||||
if modules_to_load is None:
|
||||
self.restore_from_path(
|
||||
path=multi_rl_module_ckpt_dir,
|
||||
component=COMPONENT_LEARNER + "/" + COMPONENT_RL_MODULE,
|
||||
),
|
||||
# Restore individual module IDs.
|
||||
else:
|
||||
for module_id in modules_to_load:
|
||||
path = multi_rl_module_ckpt_dir + "/" + module_id
|
||||
self.restore_from_path(
|
||||
path=path,
|
||||
component=(
|
||||
COMPONENT_LEARNER
|
||||
+ "/"
|
||||
+ COMPONENT_RL_MODULE
|
||||
+ "/"
|
||||
+ module_id
|
||||
),
|
||||
)
|
||||
if rl_module_ckpt_dirs:
|
||||
for module_id, path in rl_module_ckpt_dirs.items():
|
||||
self.restore_from_path(
|
||||
path=path,
|
||||
component=(
|
||||
COMPONENT_LEARNER + "/" + COMPONENT_RL_MODULE + "/" + module_id
|
||||
),
|
||||
)
|
||||
@@ -0,0 +1,340 @@
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
import gymnasium as gym
|
||||
import numpy as np
|
||||
|
||||
import ray
|
||||
from ray.rllib.core import DEFAULT_MODULE_ID
|
||||
from ray.rllib.core.learner.learner import Learner
|
||||
from ray.rllib.core.testing.testing_learner import BaseTestingAlgorithmConfig
|
||||
from ray.rllib.policy.sample_batch import MultiAgentBatch
|
||||
from ray.rllib.utils.framework import try_import_torch
|
||||
from ray.rllib.utils.metrics import (
|
||||
ALL_MODULES,
|
||||
MODULE_TRAIN_BATCH_SIZE_MEAN,
|
||||
NUM_ENV_STEPS_TRAINED,
|
||||
NUM_ENV_STEPS_TRAINED_LIFETIME,
|
||||
NUM_MODULE_STEPS_TRAINED,
|
||||
NUM_MODULE_STEPS_TRAINED_LIFETIME,
|
||||
WEIGHTS_SEQ_NO,
|
||||
)
|
||||
from ray.rllib.utils.numpy import convert_to_numpy
|
||||
from ray.rllib.utils.test_utils import check, get_cartpole_dataset_reader
|
||||
|
||||
torch, _ = try_import_torch()
|
||||
|
||||
|
||||
class TestLearner(unittest.TestCase):
|
||||
|
||||
ENV = gym.make("CartPole-v1")
|
||||
|
||||
@classmethod
|
||||
def setUp(cls) -> None:
|
||||
ray.init()
|
||||
|
||||
@classmethod
|
||||
def tearDown(cls) -> None:
|
||||
ray.shutdown()
|
||||
|
||||
def test_end_to_end_update(self):
|
||||
"""Tests the end-to-end update process for a single-agent scenario.
|
||||
|
||||
We check that the loss is decreasing and that the metrics are where we expect them and that values are as expected.
|
||||
"""
|
||||
|
||||
config = BaseTestingAlgorithmConfig()
|
||||
|
||||
learner = config.build_learner(env=self.ENV)
|
||||
reader = get_cartpole_dataset_reader(batch_size=512)
|
||||
|
||||
for seq_num in range(1, 1000):
|
||||
batch = reader.next().as_multi_agent()
|
||||
batch = learner._convert_batch_type(batch)
|
||||
results = learner.update(batch=batch)
|
||||
self.assertEqual(
|
||||
batch.count, results[DEFAULT_MODULE_ID][MODULE_TRAIN_BATCH_SIZE_MEAN]
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
batch.count, results[DEFAULT_MODULE_ID][NUM_MODULE_STEPS_TRAINED]
|
||||
)
|
||||
self.assertEqual(
|
||||
batch.count,
|
||||
results[DEFAULT_MODULE_ID][NUM_MODULE_STEPS_TRAINED_LIFETIME],
|
||||
)
|
||||
self.assertEqual(seq_num, results[DEFAULT_MODULE_ID][WEIGHTS_SEQ_NO])
|
||||
self.assertEqual(
|
||||
batch.count, results[DEFAULT_MODULE_ID][MODULE_TRAIN_BATCH_SIZE_MEAN]
|
||||
)
|
||||
self.assertTrue(learner.TOTAL_LOSS_KEY in results[DEFAULT_MODULE_ID])
|
||||
self.assertEqual(
|
||||
batch.count, results[ALL_MODULES][NUM_MODULE_STEPS_TRAINED]
|
||||
)
|
||||
self.assertEqual(
|
||||
batch.count, results[ALL_MODULES][NUM_MODULE_STEPS_TRAINED_LIFETIME]
|
||||
)
|
||||
self.assertEqual(batch.count, results[ALL_MODULES][NUM_ENV_STEPS_TRAINED])
|
||||
self.assertEqual(
|
||||
batch.count, results[ALL_MODULES][NUM_ENV_STEPS_TRAINED_LIFETIME]
|
||||
)
|
||||
|
||||
self.assertLess(results[DEFAULT_MODULE_ID][Learner.TOTAL_LOSS_KEY], 0.58)
|
||||
|
||||
def test_compute_gradients(self):
|
||||
"""Tests the compute_gradients correctness.
|
||||
|
||||
Tests that if we sum all the trainable variables the gradient of output w.r.t.
|
||||
the weights is all ones.
|
||||
"""
|
||||
config = BaseTestingAlgorithmConfig()
|
||||
|
||||
learner = config.build_learner(env=self.ENV)
|
||||
|
||||
params = learner.get_parameters(learner.module[DEFAULT_MODULE_ID])
|
||||
|
||||
tape = None
|
||||
loss_per_module = {ALL_MODULES: sum(param.sum() for param in params)}
|
||||
|
||||
gradients = learner.compute_gradients(loss_per_module, gradient_tape=tape)
|
||||
|
||||
# Type should be a mapping from ParamRefs to gradients.
|
||||
self.assertIsInstance(gradients, dict)
|
||||
|
||||
for grad in gradients.values():
|
||||
check(grad, np.ones(grad.shape))
|
||||
|
||||
def test_postprocess_gradients(self):
|
||||
"""Tests the base grad clipping logic in `postprocess_gradients()`."""
|
||||
|
||||
# Clip by value only.
|
||||
config = BaseTestingAlgorithmConfig().training(
|
||||
lr=0.0003, grad_clip=0.75, grad_clip_by="value"
|
||||
)
|
||||
|
||||
learner = config.build_learner(env=self.ENV)
|
||||
# Pretend our computed gradients are our weights + 1.0.
|
||||
grads = {
|
||||
learner.get_param_ref(v): v + 1.0
|
||||
for v in learner.get_parameters(learner.module[DEFAULT_MODULE_ID])
|
||||
}
|
||||
# Call the learner's postprocessing method.
|
||||
processed_grads = list(learner.postprocess_gradients(grads).values())
|
||||
# Check clipped gradients.
|
||||
# No single gradient must be larger than 0.1 or smaller than -0.1:
|
||||
self.assertTrue(
|
||||
all(
|
||||
np.max(grad) <= config.grad_clip and np.min(grad) >= -config.grad_clip
|
||||
for grad in convert_to_numpy(processed_grads)
|
||||
)
|
||||
)
|
||||
|
||||
# Clip by norm.
|
||||
config.grad_clip = 1.0
|
||||
config.grad_clip_by = "norm"
|
||||
learner = config.build_learner(env=self.ENV)
|
||||
# Pretend our computed gradients are our weights + 1.0.
|
||||
grads = {
|
||||
learner.get_param_ref(v): v + 1.0
|
||||
for v in learner.get_parameters(learner.module[DEFAULT_MODULE_ID])
|
||||
}
|
||||
# Call the learner's postprocessing method.
|
||||
processed_grads = list(learner.postprocess_gradients(grads).values())
|
||||
# Check clipped gradients.
|
||||
for proc_grad, grad in zip(
|
||||
convert_to_numpy(processed_grads),
|
||||
convert_to_numpy(list(grads.values())),
|
||||
):
|
||||
l2_norm = np.sqrt(np.sum(grad**2.0))
|
||||
if l2_norm > config.grad_clip:
|
||||
check(proc_grad, grad * (config.grad_clip / l2_norm))
|
||||
|
||||
# Clip by global norm.
|
||||
config.grad_clip = 5.0
|
||||
config.grad_clip_by = "global_norm"
|
||||
learner = config.build_learner(env=self.ENV)
|
||||
# Pretend our computed gradients are our weights + 1.0.
|
||||
grads = {
|
||||
learner.get_param_ref(v): v + 1.0
|
||||
for v in learner.get_parameters(learner.module[DEFAULT_MODULE_ID])
|
||||
}
|
||||
# Call the learner's postprocessing method.
|
||||
processed_grads = list(learner.postprocess_gradients(grads).values())
|
||||
# Check clipped gradients.
|
||||
global_norm = np.sqrt(
|
||||
np.sum(
|
||||
[np.sum(grad**2.0) for grad in convert_to_numpy(list(grads.values()))]
|
||||
)
|
||||
)
|
||||
if global_norm > config.grad_clip:
|
||||
for proc_grad, grad in zip(
|
||||
convert_to_numpy(processed_grads),
|
||||
grads.values(),
|
||||
):
|
||||
check(proc_grad, grad * (config.grad_clip / global_norm))
|
||||
|
||||
def test_apply_gradients(self):
|
||||
"""Tests the apply_gradients correctness.
|
||||
|
||||
Tests that if we apply gradients of all ones, the new params are equal to the
|
||||
standard SGD/Adam update rule.
|
||||
"""
|
||||
config = BaseTestingAlgorithmConfig().training(lr=0.0003)
|
||||
|
||||
learner = config.build_learner(env=self.ENV)
|
||||
|
||||
# calculated the expected new params based on gradients of all ones.
|
||||
params = learner.get_parameters(learner.module[DEFAULT_MODULE_ID])
|
||||
n_steps = 100
|
||||
expected = [
|
||||
(
|
||||
convert_to_numpy(param)
|
||||
- n_steps * learner.config.lr * np.ones(param.shape)
|
||||
)
|
||||
for param in params
|
||||
]
|
||||
for _ in range(n_steps):
|
||||
gradients = {learner.get_param_ref(p): torch.ones_like(p) for p in params}
|
||||
learner.apply_gradients(gradients)
|
||||
|
||||
check(params, expected)
|
||||
|
||||
def test_add_remove_module(self):
|
||||
"""Tests the compute/apply_gradients with add/remove modules.
|
||||
|
||||
Tests that if we add a module with SGD optimizer with a known lr (different
|
||||
from default), and remove the default module, with a loss that is the sum of
|
||||
all variables the updated parameters follow the SGD update rule.
|
||||
"""
|
||||
config = BaseTestingAlgorithmConfig().training(lr=0.0003)
|
||||
|
||||
learner = config.build_learner(env=self.ENV)
|
||||
rl_module_spec = config.get_default_rl_module_spec()
|
||||
rl_module_spec.observation_space = self.ENV.observation_space
|
||||
rl_module_spec.action_space = self.ENV.action_space
|
||||
learner.add_module(
|
||||
module_id="test",
|
||||
module_spec=rl_module_spec,
|
||||
)
|
||||
learner.remove_module(DEFAULT_MODULE_ID)
|
||||
|
||||
# only test module should be left
|
||||
self.assertEqual(set(learner.module.keys()), {"test"})
|
||||
|
||||
# calculated the expected new params based on gradients of all ones.
|
||||
params = learner.get_parameters(learner.module["test"])
|
||||
n_steps = 100
|
||||
expected = [
|
||||
convert_to_numpy(param) - n_steps * learner.config.lr * np.ones(param.shape)
|
||||
for param in params
|
||||
]
|
||||
for _ in range(n_steps):
|
||||
tape = None
|
||||
loss_per_module = {ALL_MODULES: sum(param.sum() for param in params)}
|
||||
gradients = learner.compute_gradients(loss_per_module, gradient_tape=tape)
|
||||
learner.apply_gradients(gradients)
|
||||
|
||||
check(params, expected)
|
||||
|
||||
def test_save_to_path_and_restore_from_path(self):
|
||||
"""Tests, whether a Learner's state is properly saved and restored."""
|
||||
config = BaseTestingAlgorithmConfig()
|
||||
|
||||
# Get a Learner instance for the framework and env.
|
||||
learner1 = config.build_learner(env=self.ENV)
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
learner1.save_to_path(tmpdir)
|
||||
|
||||
learner2 = config.build_learner(env=self.ENV)
|
||||
learner2.restore_from_path(tmpdir)
|
||||
self._check_learner_states("torch", learner1, learner2)
|
||||
|
||||
# Add a module then save/load and check states.
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
rl_module_spec = config.get_default_rl_module_spec()
|
||||
rl_module_spec.observation_space = self.ENV.observation_space
|
||||
rl_module_spec.action_space = self.ENV.action_space
|
||||
learner1.add_module(
|
||||
module_id="test",
|
||||
module_spec=rl_module_spec,
|
||||
)
|
||||
learner1.save_to_path(tmpdir)
|
||||
learner2 = Learner.from_checkpoint(tmpdir)
|
||||
self._check_learner_states("torch", learner1, learner2)
|
||||
|
||||
# Remove a module then save/load and check states.
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
learner1.remove_module(module_id=DEFAULT_MODULE_ID)
|
||||
learner1.save_to_path(tmpdir)
|
||||
learner2 = Learner.from_checkpoint(tmpdir)
|
||||
self._check_learner_states("torch", learner1, learner2)
|
||||
|
||||
def _check_learner_states(self, framework, learner1, learner2):
|
||||
check(learner1.module.get_state(), learner2.module.get_state())
|
||||
check(learner1._get_optimizer_state(), learner2._get_optimizer_state())
|
||||
check(learner1._module_optimizers, learner2._module_optimizers)
|
||||
|
||||
def test_multi_agent_learner_results(self):
|
||||
"""Tests the learner results for a multi-agent scenario.
|
||||
|
||||
We check that all metrics are where we expect them and that values are as expected.
|
||||
"""
|
||||
config = BaseTestingAlgorithmConfig()
|
||||
|
||||
learner = config.build_learner(env=self.ENV)
|
||||
learner.remove_module(module_id=DEFAULT_MODULE_ID)
|
||||
learner.add_module(
|
||||
module_id="mod1", module_spec=config.get_rl_module_spec(env=self.ENV)
|
||||
)
|
||||
learner.add_module(
|
||||
module_id="mod2", module_spec=config.get_rl_module_spec(env=self.ENV)
|
||||
)
|
||||
reader = get_cartpole_dataset_reader(batch_size=512)
|
||||
|
||||
results = {}
|
||||
for seq_num in range(1, 5):
|
||||
batch1 = reader.next()
|
||||
batch2 = reader.next()
|
||||
multi_agent_batch = MultiAgentBatch(
|
||||
{"mod1": batch1, "mod2": batch2}, batch1.count + batch2.count
|
||||
)
|
||||
batch = learner._convert_batch_type(multi_agent_batch)
|
||||
results = learner.update(batch)
|
||||
# Lifetime steps are aggregated at the root, so the return value in the results will contain only the last step.
|
||||
for module_id, sa_batch_count in zip(
|
||||
["mod1", "mod2"], [batch1.count, batch2.count]
|
||||
):
|
||||
self.assertEqual(
|
||||
sa_batch_count,
|
||||
results[module_id][NUM_MODULE_STEPS_TRAINED_LIFETIME],
|
||||
)
|
||||
self.assertEqual(seq_num, results[module_id][WEIGHTS_SEQ_NO])
|
||||
self.assertEqual(
|
||||
sa_batch_count, results[module_id][MODULE_TRAIN_BATCH_SIZE_MEAN]
|
||||
)
|
||||
# We don't know what the value should be, just check for existence.
|
||||
self.assertTrue(learner.TOTAL_LOSS_KEY in results[module_id])
|
||||
|
||||
self.assertEqual(
|
||||
batch1.count + batch2.count,
|
||||
results[ALL_MODULES][NUM_MODULE_STEPS_TRAINED_LIFETIME],
|
||||
)
|
||||
self.assertEqual(
|
||||
batch1.count + batch2.count,
|
||||
results[ALL_MODULES][NUM_MODULE_STEPS_TRAINED],
|
||||
)
|
||||
self.assertEqual(
|
||||
batch1.count + batch2.count,
|
||||
results[ALL_MODULES][NUM_ENV_STEPS_TRAINED_LIFETIME],
|
||||
)
|
||||
self.assertEqual(
|
||||
batch1.count + batch2.count, results[ALL_MODULES][NUM_ENV_STEPS_TRAINED]
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,573 @@
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
import gymnasium as gym
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
|
||||
from ray.rllib.algorithms.bc import BCConfig
|
||||
from ray.rllib.core import (
|
||||
COMPONENT_LEARNER,
|
||||
COMPONENT_RL_MODULE,
|
||||
DEFAULT_MODULE_ID,
|
||||
Columns,
|
||||
)
|
||||
from ray.rllib.core.learner.learner import Learner
|
||||
from ray.rllib.core.rl_module.multi_rl_module import MultiRLModule, MultiRLModuleSpec
|
||||
from ray.rllib.core.rl_module.rl_module import RLModuleSpec
|
||||
from ray.rllib.core.testing.testing_learner import BaseTestingAlgorithmConfig
|
||||
from ray.rllib.core.testing.torch.bc_learner import BCTorchLearner
|
||||
from ray.rllib.core.testing.torch.bc_module import DiscreteBCTorchModule
|
||||
from ray.rllib.env.multi_agent_episode import MultiAgentEpisode
|
||||
from ray.rllib.env.single_agent_episode import SingleAgentEpisode
|
||||
from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole
|
||||
from ray.rllib.utils.metrics import ALL_MODULES, LEARNER_CONNECTOR
|
||||
from ray.rllib.utils.metrics.metrics_logger import MetricsLogger
|
||||
from ray.rllib.utils.test_utils import check
|
||||
from ray.util.timer import _Timer
|
||||
|
||||
REMOTE_CONFIGS = {
|
||||
"remote-cpu": AlgorithmConfig.overrides(num_learners=1),
|
||||
"remote-gpu": AlgorithmConfig.overrides(num_learners=1, num_gpus_per_learner=1),
|
||||
"multi-gpu-ddp": AlgorithmConfig.overrides(num_learners=2, num_gpus_per_learner=1),
|
||||
"multi-cpu-ddp": AlgorithmConfig.overrides(num_learners=2, num_cpus_per_learner=2),
|
||||
# "multi-gpu-ddp-pipeline": AlgorithmConfig.overrides(
|
||||
# num_learners=2, num_gpus_per_learner=2
|
||||
# ),
|
||||
}
|
||||
|
||||
|
||||
LOCAL_CONFIGS = {
|
||||
"local-cpu": AlgorithmConfig.overrides(num_learners=0, num_gpus_per_learner=0),
|
||||
"local-gpu": AlgorithmConfig.overrides(num_learners=0, num_gpus_per_learner=1),
|
||||
}
|
||||
|
||||
|
||||
FAKE_EPISODES = [
|
||||
SingleAgentEpisode(
|
||||
observation_space=gym.spaces.Box(-1.0, 1.0, (4,), np.float32),
|
||||
observations=[
|
||||
np.array([0.1, 0.2, 0.3, 0.4], dtype=np.float32),
|
||||
np.array([0.5, 0.6, 0.7, 0.8], dtype=np.float32),
|
||||
np.array([0.9, 1.0, 1.1, 1.2], dtype=np.float32),
|
||||
np.array([0.1, 0.2, 0.3, 0.4], dtype=np.float32),
|
||||
np.array([-0.1, -0.2, -0.3, -0.4], dtype=np.float32),
|
||||
],
|
||||
action_space=gym.spaces.Discrete(2),
|
||||
actions=[0, 1, 1, 0],
|
||||
rewards=[1.0, -1.0, 0.5, 0.3],
|
||||
terminated=True,
|
||||
len_lookback_buffer=0, # all data part of actual episode
|
||||
),
|
||||
]
|
||||
|
||||
FAKE_MA_EPISODES = [
|
||||
MultiAgentEpisode(
|
||||
agent_module_ids={
|
||||
0: "p0",
|
||||
1: "p1",
|
||||
},
|
||||
observation_space=gym.spaces.Dict(
|
||||
{
|
||||
0: FAKE_EPISODES[0].observation_space,
|
||||
1: FAKE_EPISODES[0].observation_space,
|
||||
}
|
||||
),
|
||||
observations=[
|
||||
{
|
||||
0: FAKE_EPISODES[0].get_observations(i),
|
||||
1: FAKE_EPISODES[0].get_observations(i),
|
||||
}
|
||||
for i in range(5)
|
||||
],
|
||||
action_space=gym.spaces.Dict(
|
||||
{
|
||||
0: FAKE_EPISODES[0].action_space,
|
||||
1: FAKE_EPISODES[0].action_space,
|
||||
}
|
||||
),
|
||||
actions=[
|
||||
{
|
||||
0: FAKE_EPISODES[0].get_actions(i),
|
||||
1: FAKE_EPISODES[0].get_actions(i),
|
||||
}
|
||||
for i in range(4)
|
||||
],
|
||||
rewards=[
|
||||
{
|
||||
0: FAKE_EPISODES[0].get_rewards(i),
|
||||
1: FAKE_EPISODES[0].get_rewards(i),
|
||||
}
|
||||
for i in range(4)
|
||||
],
|
||||
len_lookback_buffer=0, # all data part of actual episode
|
||||
),
|
||||
]
|
||||
FAKE_MA_EPISODES[0].to_numpy()
|
||||
|
||||
FAKE_MA_EPISODES_WO_P1 = [
|
||||
MultiAgentEpisode(
|
||||
agent_module_ids={0: "p0"},
|
||||
observation_space=gym.spaces.Dict({0: FAKE_EPISODES[0].observation_space}),
|
||||
observations=[{0: FAKE_EPISODES[0].get_observations(i)} for i in range(5)],
|
||||
action_space=gym.spaces.Dict({0: FAKE_EPISODES[0].action_space}),
|
||||
actions=[{0: FAKE_EPISODES[0].get_actions(i)} for i in range(4)],
|
||||
rewards=[{0: FAKE_EPISODES[0].get_rewards(i)} for i in range(4)],
|
||||
len_lookback_buffer=0, # all data part of actual episode
|
||||
),
|
||||
]
|
||||
FAKE_MA_EPISODES_WO_P1[0].to_numpy()
|
||||
|
||||
|
||||
class TestLearnerGroupSyncUpdate(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls) -> None:
|
||||
ray.init()
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls) -> None:
|
||||
ray.shutdown()
|
||||
|
||||
def test_learner_group_build_from_algorithm_config(self):
|
||||
"""Tests whether we can build a learner_groupobject from algorithm_config."""
|
||||
|
||||
env = gym.make("CartPole-v1")
|
||||
|
||||
# Config that has its own learner class and RLModule spec.
|
||||
config = BaseTestingAlgorithmConfig()
|
||||
learner_group = config.build_learner_group(env=env)
|
||||
print(learner_group)
|
||||
learner_group.shutdown()
|
||||
|
||||
# Config for which user defines custom learner class and RLModule spec.
|
||||
config = (
|
||||
BaseTestingAlgorithmConfig()
|
||||
.training(learner_class=BCTorchLearner)
|
||||
.rl_module(
|
||||
rl_module_spec=RLModuleSpec(
|
||||
module_class=DiscreteBCTorchModule,
|
||||
model_config={"fcnet_hiddens": [32]},
|
||||
)
|
||||
)
|
||||
)
|
||||
learner_group = config.build_learner_group(env=env)
|
||||
print(learner_group)
|
||||
learner_group.shutdown()
|
||||
|
||||
def test_update_multi_gpu(self):
|
||||
return
|
||||
|
||||
scaling_modes = ["multi-gpu-ddp", "remote-gpu"]
|
||||
|
||||
for scaling_mode in scaling_modes:
|
||||
print(f"Testing scaling mode: {scaling_mode}.")
|
||||
env = gym.make("CartPole-v1")
|
||||
|
||||
config_overrides = REMOTE_CONFIGS[scaling_mode]
|
||||
config = BaseTestingAlgorithmConfig().update_from_dict(config_overrides)
|
||||
learner_group = config.build_learner_group(env=env)
|
||||
|
||||
min_loss = float("inf")
|
||||
for iter_i in range(1000):
|
||||
results = learner_group.update(episodes=FAKE_EPISODES)
|
||||
|
||||
loss = np.mean(
|
||||
[res[ALL_MODULES][Learner.TOTAL_LOSS_KEY] for res in results]
|
||||
)
|
||||
min_loss = min(loss, min_loss)
|
||||
print(f"[iter = {iter_i}] Loss: {loss:.3f}, Min Loss: {min_loss:.3f}")
|
||||
# The loss is initially around 0.69 (ln2). When it gets to around
|
||||
# 0.57 the return of the policy gets to around 100.
|
||||
if min_loss < 0.57:
|
||||
break
|
||||
|
||||
for res1, res2 in zip(results, results[1:]):
|
||||
self.assertEqual(
|
||||
res1[DEFAULT_MODULE_ID]["mean_weight"],
|
||||
res2[DEFAULT_MODULE_ID]["mean_weight"],
|
||||
)
|
||||
|
||||
self.assertLess(min_loss, 0.57)
|
||||
|
||||
# Make sure the learner_group resources are freed up so that we don't
|
||||
# autoscale.
|
||||
learner_group.shutdown()
|
||||
del learner_group
|
||||
|
||||
def test_add_module_and_remove_module(self):
|
||||
scaling_modes = ["local-cpu", "multi-cpu-ddp"]
|
||||
|
||||
for scaling_mode in scaling_modes:
|
||||
print(f"Testing scaling mode: {scaling_mode}.")
|
||||
ma_env = MultiAgentCartPole({"num_agents": 2})
|
||||
config_overrides = REMOTE_CONFIGS.get(scaling_mode) or LOCAL_CONFIGS.get(
|
||||
scaling_mode
|
||||
)
|
||||
config = (
|
||||
BCConfig()
|
||||
.update_from_dict(config_overrides)
|
||||
.multi_agent(
|
||||
policies={"p0"},
|
||||
policy_mapping_fn=lambda aid, *ar, **kw: f"p{aid}",
|
||||
)
|
||||
.rl_module(
|
||||
rl_module_spec=MultiRLModuleSpec(
|
||||
rl_module_specs={"p0": RLModuleSpec()},
|
||||
)
|
||||
)
|
||||
)
|
||||
learner_group = config.build_learner_group(env=ma_env)
|
||||
|
||||
# Update once with the default policy.
|
||||
learner_group.update(episodes=FAKE_MA_EPISODES_WO_P1)
|
||||
|
||||
# Add a test_module.
|
||||
learner_group.add_module(
|
||||
module_id="p1",
|
||||
module_spec=config.get_multi_rl_module_spec(env=ma_env).module_specs[
|
||||
"p0"
|
||||
],
|
||||
)
|
||||
# Do training that includes the test_module.
|
||||
results = learner_group.update(episodes=FAKE_MA_EPISODES)
|
||||
|
||||
# check that module ids are updated to include the new module
|
||||
module_ids_after_add = {"p0", "p1"}
|
||||
# Compare module IDs in results with expected ones.
|
||||
self.assertEqual(
|
||||
set(results[0].keys()) - {ALL_MODULES}, module_ids_after_add
|
||||
)
|
||||
|
||||
# Remove the test_module.
|
||||
learner_group.remove_module(module_id="p1")
|
||||
|
||||
# Run training without the test_module.
|
||||
results = learner_group.update(episodes=FAKE_MA_EPISODES_WO_P1)
|
||||
|
||||
# check that module ids are updated after remove operation to not
|
||||
# include the new module
|
||||
# remove the total_loss key since its not a module key
|
||||
self.assertEqual(set(results[0].keys()) - {ALL_MODULES}, {"p0"})
|
||||
|
||||
# make sure the learner_group resources are freed up so that we don't
|
||||
# autoscale
|
||||
learner_group.shutdown()
|
||||
del learner_group
|
||||
|
||||
|
||||
class TestLearnerGroupCheckpointRestore(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls) -> None:
|
||||
ray.init()
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls) -> None:
|
||||
ray.shutdown()
|
||||
|
||||
def test_restore_from_path_multi_rl_module_and_individual_modules(self):
|
||||
"""Tests whether MultiRLModule- and single RLModule states can be restored."""
|
||||
# this is expanded to more scaling modes on the release ci.
|
||||
scaling_modes = ["local-cpu", "multi-gpu-ddp"]
|
||||
|
||||
for scaling_mode in scaling_modes:
|
||||
print(f"Testing scaling mode: {scaling_mode}.")
|
||||
# env will have agent ids 0 and 1
|
||||
env = MultiAgentCartPole({"num_agents": 2})
|
||||
|
||||
config_overrides = REMOTE_CONFIGS.get(scaling_mode) or LOCAL_CONFIGS.get(
|
||||
scaling_mode
|
||||
)
|
||||
config = BaseTestingAlgorithmConfig().update_from_dict(config_overrides)
|
||||
learner_group = config.build_learner_group(env=env)
|
||||
spec = config.get_multi_rl_module_spec(env=env).module_specs[
|
||||
DEFAULT_MODULE_ID
|
||||
]
|
||||
learner_group.add_module(module_id="0", module_spec=spec)
|
||||
learner_group.add_module(module_id="1", module_spec=spec)
|
||||
learner_group.remove_module(DEFAULT_MODULE_ID)
|
||||
|
||||
module_0 = spec.build()
|
||||
module_1 = spec.build()
|
||||
multi_rl_module = MultiRLModule()
|
||||
multi_rl_module.add_module(module_id="0", module=module_0)
|
||||
multi_rl_module.add_module(module_id="1", module=module_1)
|
||||
|
||||
# Check if we can load just the MultiRLModule.
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
multi_rl_module.save_to_path(tmpdir)
|
||||
old_learner_weights = learner_group.get_weights()
|
||||
learner_group.restore_from_path(
|
||||
tmpdir,
|
||||
component=COMPONENT_LEARNER + "/" + COMPONENT_RL_MODULE,
|
||||
)
|
||||
# Check the weights of the module in the learner group are the
|
||||
# same as the weights of the newly created MultiRLModule
|
||||
check(learner_group.get_weights(), multi_rl_module.get_state())
|
||||
learner_group.set_state(
|
||||
{
|
||||
COMPONENT_LEARNER: {COMPONENT_RL_MODULE: old_learner_weights},
|
||||
}
|
||||
)
|
||||
check(learner_group.get_weights(), old_learner_weights)
|
||||
|
||||
# Check if we can load just single agent RL Modules.
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
module_0.save_to_path(tmpdir)
|
||||
with tempfile.TemporaryDirectory() as tmpdir2:
|
||||
temp_module = spec.build()
|
||||
temp_module.save_to_path(tmpdir2)
|
||||
|
||||
old_learner_weights = learner_group.get_weights()
|
||||
learner_group.restore_from_path(
|
||||
tmpdir,
|
||||
component=COMPONENT_LEARNER + "/" + COMPONENT_RL_MODULE + "/0",
|
||||
)
|
||||
learner_group.restore_from_path(
|
||||
tmpdir2,
|
||||
component=COMPONENT_LEARNER + "/" + COMPONENT_RL_MODULE + "/1",
|
||||
)
|
||||
# check the weights of the module in the learner group are the
|
||||
# same as the weights of the newly created MultiRLModule
|
||||
new_multi_rl_module = MultiRLModule()
|
||||
new_multi_rl_module.add_module(module_id="0", module=module_0)
|
||||
new_multi_rl_module.add_module(module_id="1", module=temp_module)
|
||||
check(learner_group.get_weights(), new_multi_rl_module.get_state())
|
||||
learner_group.set_weights(old_learner_weights)
|
||||
|
||||
# Check if we can first load a MultiRLModule, then a single agent RLModule
|
||||
# (within that MultiRLModule). Check that the single agent RL Module is
|
||||
# loaded over the matching submodule in the MultiRLModule.
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
module_0 = spec.build()
|
||||
multi_rl_module = MultiRLModule()
|
||||
multi_rl_module.add_module(module_id="0", module=module_0)
|
||||
multi_rl_module.add_module(module_id="1", module=spec.build())
|
||||
multi_rl_module.save_to_path(tmpdir)
|
||||
with tempfile.TemporaryDirectory() as tmpdir2:
|
||||
module_1 = spec.build()
|
||||
module_1.save_to_path(tmpdir2)
|
||||
learner_group.restore_from_path(
|
||||
tmpdir,
|
||||
component=COMPONENT_LEARNER + "/" + COMPONENT_RL_MODULE,
|
||||
)
|
||||
learner_group.restore_from_path(
|
||||
tmpdir2,
|
||||
component=COMPONENT_LEARNER + "/" + COMPONENT_RL_MODULE + "/1",
|
||||
)
|
||||
new_multi_rl_module = MultiRLModule()
|
||||
new_multi_rl_module.add_module(module_id="0", module=module_0)
|
||||
new_multi_rl_module.add_module(module_id="1", module=module_1)
|
||||
check(learner_group.get_weights(), new_multi_rl_module.get_state())
|
||||
del learner_group
|
||||
|
||||
|
||||
class TestLearnerGroupSaveAndRestoreState(unittest.TestCase):
|
||||
|
||||
FAKE_BATCH = {
|
||||
Columns.OBS: np.array(
|
||||
[
|
||||
[0.1, 0.2, 0.3, 0.4],
|
||||
[0.5, 0.6, 0.7, 0.8],
|
||||
[0.9, 1.0, 1.1, 1.2],
|
||||
[1.3, 1.4, 1.5, 1.6],
|
||||
],
|
||||
dtype=np.float32,
|
||||
),
|
||||
Columns.NEXT_OBS: np.array(
|
||||
[
|
||||
[0.1, 0.2, 0.3, 0.4],
|
||||
[0.5, 0.6, 0.7, 0.8],
|
||||
[0.9, 1.0, 1.1, 1.2],
|
||||
[1.3, 1.4, 1.5, 1.6],
|
||||
],
|
||||
dtype=np.float32,
|
||||
),
|
||||
Columns.ACTIONS: np.array([0, 1, 1, 0]),
|
||||
Columns.REWARDS: np.array([1.0, -1.0, 0.5, 0.6], dtype=np.float32),
|
||||
Columns.TERMINATEDS: np.array([False, False, True, False]),
|
||||
Columns.TRUNCATEDS: np.array([False, False, False, False]),
|
||||
Columns.VF_PREDS: np.array([0.5, 0.6, 0.7, 0.8], dtype=np.float32),
|
||||
Columns.ACTION_DIST_INPUTS: np.array(
|
||||
[[-2.0, 0.5], [-3.0, -0.3], [-0.1, 2.5], [-0.2, 3.5]], dtype=np.float32
|
||||
),
|
||||
Columns.ACTION_LOGP: np.array([-0.5, -0.1, -0.2, -0.3], dtype=np.float32),
|
||||
Columns.EPS_ID: np.array([0, 0, 0, 0]),
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls) -> None:
|
||||
ray.init()
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls) -> None:
|
||||
ray.shutdown()
|
||||
|
||||
def test_save_to_path_and_restore_from_path(self):
|
||||
"""Check that saving and loading learner group state works."""
|
||||
# this is expanded to more scaling modes on the release ci.
|
||||
scaling_modes = ["local-cpu"] # , "multi-gpu-ddp"]
|
||||
|
||||
for scaling_mode in scaling_modes:
|
||||
print(f"Testing scaling mode: {scaling_mode}.")
|
||||
env = gym.make("CartPole-v1")
|
||||
|
||||
config_overrides = REMOTE_CONFIGS.get(scaling_mode) or LOCAL_CONFIGS.get(
|
||||
scaling_mode
|
||||
)
|
||||
config = BaseTestingAlgorithmConfig().update_from_dict(config_overrides)
|
||||
learner_group = config.build_learner_group(env=env)
|
||||
|
||||
# Checkpoint the initial learner state for later comparison.
|
||||
initial_learner_checkpoint_dir = tempfile.TemporaryDirectory().name
|
||||
learner_group.save_to_path(initial_learner_checkpoint_dir)
|
||||
# Test the convenience method `.get_weights()`.
|
||||
initial_weights = learner_group.get_weights()
|
||||
|
||||
# Do a single update.
|
||||
learner_group.update(episodes=FAKE_EPISODES)
|
||||
weights_after_update = learner_group.get_state(
|
||||
components=COMPONENT_LEARNER + "/" + COMPONENT_RL_MODULE
|
||||
)[COMPONENT_LEARNER][COMPONENT_RL_MODULE]
|
||||
# Weights after the update must be different from original ones.
|
||||
check(initial_weights, weights_after_update, false=True)
|
||||
|
||||
# Checkpoint the learner state after 1 update for later comparison.
|
||||
learner_after_1_update_checkpoint_dir = tempfile.TemporaryDirectory().name
|
||||
learner_group.save_to_path(learner_after_1_update_checkpoint_dir)
|
||||
|
||||
# Remove that learner, construct a new one, and load the state of the old
|
||||
# learner into the new one.
|
||||
learner_group.shutdown()
|
||||
del learner_group
|
||||
|
||||
learner_group = config.build_learner_group(env=env)
|
||||
learner_group.restore_from_path(learner_after_1_update_checkpoint_dir)
|
||||
|
||||
# Do another update.
|
||||
results_2nd_update_with_break = learner_group.update(episodes=FAKE_EPISODES)
|
||||
weights_after_2_updates_with_break = learner_group.get_state(
|
||||
components=COMPONENT_LEARNER + "/" + COMPONENT_RL_MODULE
|
||||
)[COMPONENT_LEARNER][COMPONENT_RL_MODULE]
|
||||
learner_group.shutdown()
|
||||
del learner_group
|
||||
|
||||
# Construct a new learner group and load the initial state of the learner.
|
||||
learner_group = config.build_learner_group(env=env)
|
||||
learner_group.restore_from_path(initial_learner_checkpoint_dir)
|
||||
weights_after_restore = learner_group.get_state(
|
||||
components=COMPONENT_LEARNER + "/" + COMPONENT_RL_MODULE
|
||||
)[COMPONENT_LEARNER][COMPONENT_RL_MODULE]
|
||||
check(initial_weights, weights_after_restore)
|
||||
# Perform 2 updates to get to the same state as the previous learners.
|
||||
learner_group.update(episodes=FAKE_EPISODES)
|
||||
results_2nd_update_without_break = learner_group.update(
|
||||
episodes=FAKE_EPISODES
|
||||
)
|
||||
weights_after_2_updates_without_break = learner_group.get_weights()
|
||||
learner_group.shutdown()
|
||||
del learner_group
|
||||
|
||||
# Compare the results of the two updates.
|
||||
for r1, r2 in zip(
|
||||
results_2nd_update_with_break,
|
||||
results_2nd_update_without_break,
|
||||
):
|
||||
r1[ALL_MODULES].pop(LEARNER_CONNECTOR)
|
||||
r2[ALL_MODULES].pop(LEARNER_CONNECTOR)
|
||||
check(
|
||||
MetricsLogger.peek_results(results_2nd_update_with_break),
|
||||
MetricsLogger.peek_results(results_2nd_update_without_break),
|
||||
rtol=0.05,
|
||||
)
|
||||
check(
|
||||
weights_after_2_updates_with_break,
|
||||
weights_after_2_updates_without_break,
|
||||
rtol=0.05,
|
||||
)
|
||||
|
||||
|
||||
class TestLearnerGroupAsyncUpdate(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls) -> None:
|
||||
ray.init()
|
||||
|
||||
@classmethod
|
||||
def tearDown(cls) -> None:
|
||||
ray.shutdown()
|
||||
|
||||
def test_async_update(self):
|
||||
"""Test that async style updates converge to the same result as sync."""
|
||||
scaling_modes = ["multi-gpu-ddp", "multi-cpu-ddp", "remote-gpu"]
|
||||
|
||||
for scaling_mode in scaling_modes:
|
||||
print(f"Testing scaling mode: {scaling_mode}.")
|
||||
env = gym.make("CartPole-v1")
|
||||
config_overrides = REMOTE_CONFIGS[scaling_mode]
|
||||
config = BaseTestingAlgorithmConfig().update_from_dict(config_overrides)
|
||||
learner_group = config.build_learner_group(env=env)
|
||||
timer_sync = _Timer()
|
||||
timer_async = _Timer()
|
||||
with timer_sync:
|
||||
learner_group.update(episodes=FAKE_EPISODES, async_update=False)
|
||||
with timer_async:
|
||||
result_async = learner_group.update(
|
||||
episodes=FAKE_EPISODES, async_update=True
|
||||
)
|
||||
# Ideally the first async update will return nothing, and an easy
|
||||
# way to check that is if the time for an async update call is faster
|
||||
# than the time for a sync update call.
|
||||
self.assertLess(timer_async.mean, timer_sync.mean)
|
||||
self.assertIsInstance(result_async, list)
|
||||
loss = float("inf")
|
||||
iter_i = 0
|
||||
while True:
|
||||
result_async = learner_group.update(
|
||||
episodes=FAKE_EPISODES, async_update=True
|
||||
)
|
||||
if not result_async:
|
||||
continue
|
||||
self.assertIsInstance(result_async, list)
|
||||
self.assertIsInstance(result_async[0], dict)
|
||||
# Check one async Learner result.
|
||||
loss = result_async[0][DEFAULT_MODULE_ID][Learner.TOTAL_LOSS_KEY]
|
||||
# The loss is initially around 0.69 (ln2). When it gets to around
|
||||
# 0.57 the return of the policy gets to around 100.
|
||||
if loss < 0.57:
|
||||
break
|
||||
# Compare reported "mean_weight" with actual ones.
|
||||
_check_multi_worker_weights(learner_group, result_async)
|
||||
iter_i += 1
|
||||
learner_group.shutdown()
|
||||
self.assertLess(loss, 0.57)
|
||||
|
||||
|
||||
def _check_multi_worker_weights(learner_group, results):
|
||||
# Check that module weights are updated across workers and synchronized.
|
||||
# for i in range(1, len(results)):
|
||||
|
||||
learner_1_results = results[0]
|
||||
for module_id, mod_result in learner_1_results.items():
|
||||
if module_id == ALL_MODULES:
|
||||
continue
|
||||
results = MetricsLogger.peek_results(results)
|
||||
reported_mean_weights = np.mean([r[module_id]["mean_weight"] for r in results])
|
||||
|
||||
# Compare the reported mean weights (merged across all Learner workers,
|
||||
# which all should have the same weights after updating) with the actual
|
||||
# current mean weights.
|
||||
parameters = learner_group.get_state(
|
||||
components=(
|
||||
COMPONENT_LEARNER + "/" + COMPONENT_RL_MODULE + "/" + module_id
|
||||
),
|
||||
)[COMPONENT_LEARNER][COMPONENT_RL_MODULE][module_id]
|
||||
actual_mean_weights = np.mean([w.mean() for w in parameters.values()])
|
||||
check(reported_mean_weights, actual_mean_weights, rtol=0.02)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
class_ = sys.argv[1] if len(sys.argv) > 1 else None
|
||||
sys.exit(pytest.main(["-v", __file__ + ("" if class_ is None else "::" + class_)]))
|
||||
@@ -0,0 +1,125 @@
|
||||
import itertools
|
||||
import unittest
|
||||
|
||||
import gymnasium as gym
|
||||
|
||||
import ray
|
||||
from ray.rllib.algorithms.algorithm_config import TorchCompileWhatToCompile
|
||||
from ray.rllib.core.testing.testing_learner import BaseTestingAlgorithmConfig
|
||||
from ray.rllib.policy.sample_batch import MultiAgentBatch
|
||||
from ray.rllib.utils.test_utils import get_cartpole_dataset_reader
|
||||
from ray.rllib.utils.torch_utils import _dynamo_is_available
|
||||
|
||||
|
||||
class TestLearner(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUp(cls) -> None:
|
||||
ray.init()
|
||||
|
||||
@classmethod
|
||||
def tearDown(cls) -> None:
|
||||
ray.shutdown()
|
||||
|
||||
# Todo (rllib-team): Fix for torch 2.0+
|
||||
@unittest.skip("Failing with torch >= 2.0")
|
||||
@unittest.skipIf(not _dynamo_is_available(), "torch._dynamo not available")
|
||||
def test_torch_compile(self):
|
||||
"""Test if torch.compile() can be applied and used on the learner.
|
||||
|
||||
Also tests if we can update with the compiled update method without errors.
|
||||
"""
|
||||
|
||||
env = gym.make("CartPole-v1")
|
||||
is_multi_agents = [False, True]
|
||||
what_to_compiles = [
|
||||
TorchCompileWhatToCompile.FORWARD_TRAIN,
|
||||
TorchCompileWhatToCompile.COMPLETE_UPDATE,
|
||||
]
|
||||
|
||||
for is_multi_agent, what_to_compile in itertools.product(
|
||||
is_multi_agents, what_to_compiles
|
||||
):
|
||||
print(
|
||||
f"Testing is_multi_agent={is_multi_agent},"
|
||||
f"what_to_compile={what_to_compile}"
|
||||
)
|
||||
config = BaseTestingAlgorithmConfig().framework(
|
||||
torch_compile_learner=True,
|
||||
torch_compile_learner_what_to_compile=what_to_compile,
|
||||
)
|
||||
learner = config.build_learner(env=env)
|
||||
|
||||
learner.build()
|
||||
|
||||
reader = get_cartpole_dataset_reader(batch_size=512)
|
||||
|
||||
for iter_i in range(10):
|
||||
batch = reader.next()
|
||||
learner.update(batch.as_multi_agent())
|
||||
|
||||
rl_module_spec = config.get_default_rl_module_spec()
|
||||
rl_module_spec.observation_space = env.observation_space
|
||||
rl_module_spec.action_space = env.action_space
|
||||
learner.add_module(module_id="another_module", module_spec=rl_module_spec)
|
||||
|
||||
for iter_i in range(10):
|
||||
batch = MultiAgentBatch(
|
||||
{"another_module": reader.next(), "default_policy": reader.next()},
|
||||
0,
|
||||
)
|
||||
learner.update(batch)
|
||||
|
||||
learner.remove_module(module_id="another_module")
|
||||
|
||||
# Todo (rllib-team): Fix for torch 2.0+
|
||||
@unittest.skip("Failing with torch >= 2.0")
|
||||
@unittest.skipIf(not _dynamo_is_available(), "torch._dynamo not available")
|
||||
def test_torch_compile_no_breaks(self):
|
||||
"""Tests if torch.compile() does encounter too many breaks.
|
||||
|
||||
torch.compile() should ideally not encounter any breaks when compiling the
|
||||
update method of the learner. This method tests if we encounter only a given
|
||||
number of breaks.
|
||||
"""
|
||||
|
||||
env = gym.make("CartPole-v1")
|
||||
|
||||
config = BaseTestingAlgorithmConfig().framework(torch_compile_learner=True)
|
||||
learner = config.build_learner(env=env)
|
||||
|
||||
import torch._dynamo as dynamo
|
||||
|
||||
reader = get_cartpole_dataset_reader(batch_size=512)
|
||||
|
||||
batch = reader.next().as_multi_agent()
|
||||
batch = learner._convert_batch_type(batch)
|
||||
|
||||
# The followingcall to dynamo.explain() breaks depending on the torch version.
|
||||
# It works for torch==2.0.0.
|
||||
# TODO(Artur): Fit this to to the correct torch version once it is enabled on
|
||||
# CI.
|
||||
# This is a helper method of dynamo to analyze where breaks occur.
|
||||
(
|
||||
explanation,
|
||||
out_guards,
|
||||
graphs,
|
||||
ops_per_graph,
|
||||
break_reasons,
|
||||
explanation_verbose,
|
||||
) = dynamo.explain(learner._update, batch)
|
||||
|
||||
print(explanation_verbose)
|
||||
|
||||
# There should be only one break reason - `return_value` - since inputs and
|
||||
# outputs are not checked
|
||||
# TODO(Artur): Attempt bringing breaks down to 1. (This may not be possible)
|
||||
# Note: This test is skipped on CI if torch dynamo is available.
|
||||
self.assertEqual(len(break_reasons), 3)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,82 @@
|
||||
import torch
|
||||
|
||||
from ray.rllib.algorithms.ppo import PPOConfig
|
||||
|
||||
|
||||
def test_ppo_checkpoint_restore_preserves_optimizer_state(tmp_path):
|
||||
config = (
|
||||
PPOConfig()
|
||||
.environment("CartPole-v1")
|
||||
.framework("torch")
|
||||
.env_runners(num_env_runners=0)
|
||||
.learners(num_learners=0)
|
||||
.training(
|
||||
train_batch_size=128,
|
||||
minibatch_size=32,
|
||||
num_epochs=1,
|
||||
)
|
||||
)
|
||||
|
||||
algo1 = config.build_algo()
|
||||
try:
|
||||
# Populate optimizer state before checkpointing.
|
||||
algo1.train()
|
||||
|
||||
learner1 = algo1.learner_group._learner
|
||||
optimizer1 = next(iter(learner1._named_optimizers.values()))
|
||||
original_state = optimizer1.state_dict()
|
||||
|
||||
ckpt_dir = algo1.save_to_path(tmp_path / "ckpt")
|
||||
finally:
|
||||
algo1.stop()
|
||||
|
||||
algo2 = config.build_algo()
|
||||
try:
|
||||
algo2.restore_from_path(ckpt_dir)
|
||||
|
||||
learner2 = algo2.learner_group._learner
|
||||
optimizer2 = next(iter(learner2._named_optimizers.values()))
|
||||
restored_state = optimizer2.state_dict()
|
||||
param_group = optimizer2.param_groups[0]
|
||||
|
||||
# Param-group metadata must remain Python scalars after restore.
|
||||
assert isinstance(param_group["betas"][0], float)
|
||||
assert isinstance(param_group["betas"][1], float)
|
||||
assert not torch.is_tensor(param_group["betas"][0])
|
||||
assert not torch.is_tensor(param_group["betas"][1])
|
||||
|
||||
# Restored optimizer state should match the original optimizer state.
|
||||
assert restored_state["param_groups"] == original_state["param_groups"]
|
||||
assert restored_state["state"].keys() == original_state["state"].keys()
|
||||
|
||||
for param_id, original_param_state in original_state["state"].items():
|
||||
restored_param_state = restored_state["state"][param_id]
|
||||
|
||||
assert torch.equal(
|
||||
restored_param_state["step"], original_param_state["step"]
|
||||
)
|
||||
assert torch.equal(
|
||||
restored_param_state["exp_avg"], original_param_state["exp_avg"]
|
||||
)
|
||||
assert torch.equal(
|
||||
restored_param_state["exp_avg_sq"], original_param_state["exp_avg_sq"]
|
||||
)
|
||||
|
||||
# Optimizer state buffers must still be tensors.
|
||||
state = next(iter(optimizer2.state.values()))
|
||||
assert torch.is_tensor(state["step"])
|
||||
assert torch.is_tensor(state["exp_avg"])
|
||||
assert torch.is_tensor(state["exp_avg_sq"])
|
||||
|
||||
# Training must continue after restore.
|
||||
algo2.train()
|
||||
finally:
|
||||
algo2.stop()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,374 @@
|
||||
import contextlib
|
||||
import logging
|
||||
from typing import Any, Dict, Optional, Tuple
|
||||
|
||||
from ray.rllib.algorithms.algorithm_config import (
|
||||
TorchCompileWhatToCompile,
|
||||
)
|
||||
from ray.rllib.core.columns import Columns
|
||||
from ray.rllib.core.learner.differentiable_learner import DifferentiableLearner
|
||||
from ray.rllib.core.rl_module.torch.torch_rl_module import (
|
||||
TorchCompileConfig,
|
||||
)
|
||||
from ray.rllib.policy.sample_batch import MultiAgentBatch
|
||||
from ray.rllib.utils.annotations import override
|
||||
from ray.rllib.utils.framework import try_import_torch
|
||||
from ray.rllib.utils.metrics import (
|
||||
DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY,
|
||||
WEIGHTS_SEQ_NO,
|
||||
)
|
||||
from ray.rllib.utils.torch_utils import convert_to_torch_tensor
|
||||
from ray.rllib.utils.typing import DeviceType, ModuleID, NamedParamDict, TensorType
|
||||
|
||||
logger = logging.getLogger("__name__")
|
||||
|
||||
torch, nn = try_import_torch()
|
||||
|
||||
|
||||
# TODO (simon): Maybe create again a base class `TorchLearnable`.
|
||||
class TorchDifferentiableLearner(DifferentiableLearner):
|
||||
"""A `DifferentiableLearner` class leveraging PyTorch for functional updates.
|
||||
|
||||
This class utilizes PyTorch 2.0's `func` module to perform functional
|
||||
updates on the provided parameters.
|
||||
"""
|
||||
|
||||
# Set the framework to `"torch"`.
|
||||
framework: str = "torch"
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
|
||||
# First initialize the `DifferentiableLearner` base class to set
|
||||
# the configurations and `MultiRLModule`.
|
||||
super().__init__(**kwargs)
|
||||
|
||||
# Whether to compile the RL Module of this learner. This implies that the.
|
||||
# forward_train method of the RL Module will be compiled. Further more,
|
||||
# other forward methods of the RL Module will be compiled on demand.
|
||||
# This is assumed to not happen, since other forwrad methods are not expected
|
||||
# to be used during training.
|
||||
self._torch_compile_forward_train = False
|
||||
self._torch_compile_cfg = None
|
||||
# Whether to compile the `_uncompiled_update` method of this learner. This
|
||||
# implies that everything within `_uncompiled_update` will be compiled,
|
||||
# not only the forward_train method of the RL Module.
|
||||
# Note that this is experimental.
|
||||
# Note that this requires recompiling the forward methods once we add/remove
|
||||
# RL Modules.
|
||||
self._torch_compile_complete_update = False
|
||||
if self.config.torch_compile_learner:
|
||||
if (
|
||||
self.config.torch_compile_learner_what_to_compile
|
||||
== TorchCompileWhatToCompile.COMPLETE_UPDATE
|
||||
):
|
||||
self._torch_compile_complete_update = True
|
||||
self._compiled_update_initialized = False
|
||||
else:
|
||||
self._torch_compile_forward_train = True
|
||||
|
||||
self._torch_compile_cfg = TorchCompileConfig(
|
||||
torch_dynamo_backend=self.config.torch_compile_learner_dynamo_backend,
|
||||
torch_dynamo_mode=self.config.torch_compile_learner_dynamo_mode,
|
||||
)
|
||||
|
||||
# TODO (simon): See, if we can include these without a torch optimizer.
|
||||
# self._lr_schedulers = {}
|
||||
# self._lr_scheduler_classes = None
|
||||
# if self.config._torch_lr_scheduler_classes:
|
||||
# self._lr_scheduler_classes = self.config._torch_lr_scheduler_classes
|
||||
|
||||
def _uncompiled_update(
|
||||
self,
|
||||
batch: Dict,
|
||||
params: Dict[ModuleID, NamedParamDict],
|
||||
**kwargs,
|
||||
) -> Tuple[Any, Any, Dict[ModuleID, NamedParamDict], Any]:
|
||||
"""Performs a single functional update using a batch of data.
|
||||
|
||||
This update operates on parameters passed via a functional call to the
|
||||
`MultiRLModule` and leverages PyTorch 2.0's `autograd` module. Parameters
|
||||
are not modified in-place within `self._module`; instead, updates are
|
||||
applied to the cloned parameters provided.
|
||||
|
||||
Args:
|
||||
batch: A dictionary (or `MultiAgentBatch`) containing training data for
|
||||
all modules in the `MultiRLModule` (that should be trained).
|
||||
params: A dictionary of named parameters for each module id.
|
||||
|
||||
Returns:
|
||||
A tuple consisting of:
|
||||
1) the output of a functional forward call to the RLModule using
|
||||
`params`,
|
||||
2) the `loss_per_module` dictionary mapping module IDs to individual
|
||||
loss tensors,
|
||||
3) the functionally updated parameters in the (dict) format passed in,
|
||||
4) a metrics dict mapping module IDs to metrics key/value pairs.
|
||||
|
||||
"""
|
||||
# TODO (sven): Causes weird cuda error when WandB is used.
|
||||
# Diagnosis thus far:
|
||||
# - All peek values during metrics.reduce are non-tensors.
|
||||
# - However, in impala.py::training_step(), a tensor does arrive after learner
|
||||
# group.update(), so somehow, there is still a race condition
|
||||
# possible (learner, which performs the reduce() and learner thread, which
|
||||
# performs the logging of tensors into metrics logger).
|
||||
self._compute_off_policyness(batch)
|
||||
|
||||
# Make a functional forward pass with the provided parameters.
|
||||
fwd_out = self._make_functional_call(params, batch)
|
||||
loss_per_module = self.compute_losses(fwd_out=fwd_out, batch=batch)
|
||||
# Compute gradients for the provided parameters.
|
||||
gradients = self.compute_gradients(loss_per_module, params)
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
if self.config.num_learners > 1:
|
||||
for mod in self.module.values():
|
||||
# Skip non-torch modules, b/c they may not have the `no_sync` API.
|
||||
if isinstance(mod, torch.nn.Module):
|
||||
stack.enter_context(mod.no_sync())
|
||||
# TODO (simon): See, if we need here postprocessing of gradients.
|
||||
# postprocessed_gradients = self.postprocess_gradients(gradients)
|
||||
# Make a stateless (of `params`) update of the `RLModule` parameters.
|
||||
params = self.apply_gradients(gradients, params)
|
||||
|
||||
# Deactivate tensor-mode on our MetricsLogger and collect the (tensor)
|
||||
# results.
|
||||
return fwd_out, loss_per_module, params, {}
|
||||
|
||||
# TODO (simon): Maybe make type for gradients.
|
||||
@override(DifferentiableLearner)
|
||||
def compute_gradients(
|
||||
self,
|
||||
loss_per_module: Dict[ModuleID, TensorType],
|
||||
params: Dict[ModuleID, NamedParamDict],
|
||||
**kwargs,
|
||||
) -> Dict[ModuleID, NamedParamDict]:
|
||||
"""Computes functionally gradients based on the given losses.
|
||||
|
||||
This method uses `torch.autograd.grad` to make the backward pass on the
|
||||
`MultiRLModule` which enables a functional backward pass. If a PyTorch
|
||||
optimizer is needed a differentiable one must be used (e.g. `torchopt`).
|
||||
|
||||
Args:
|
||||
loss_per_module: Dict mapping module IDs to their individual total loss
|
||||
terms, computed by the individual `compute_loss_for_module()` calls.
|
||||
The overall total loss (sum of loss terms over all modules) is stored
|
||||
under `loss_per_module[ALL_MODULES]`
|
||||
params: A dictionary containing named parameters for each module id.
|
||||
**kwargs: Forward compatibility kwargs.
|
||||
|
||||
Returns:
|
||||
The (named) gradients in the same (dict) format as `params`.
|
||||
"""
|
||||
|
||||
# TODO (simon): Add grad scalers later.
|
||||
total_loss = sum(loss_per_module.values())
|
||||
|
||||
# Use `torch`'s `autograd` to compute gradients and create a graph, so we can
|
||||
# compute higher order gradients. Allow specified inputs not being used in outputs
|
||||
# as probably not all modules/parameters of a `MultiRLModule` are used in the loss.
|
||||
# Note, parameters are named parameters as this is needed by the
|
||||
# `torch.func.functional_call`
|
||||
# TODO (simon): Make sure this works for `MultiRLModule`s. This here can have
|
||||
# all parameter tensors in a list. But the `functional_call` above needs named
|
||||
# parameters for each module. Implement this via `foreach_module`.
|
||||
grads = torch.autograd.grad(
|
||||
total_loss,
|
||||
sum((list(param.values()) for mid, param in params.items()), []),
|
||||
create_graph=True,
|
||||
retain_graph=True,
|
||||
allow_unused=True,
|
||||
)
|
||||
|
||||
# Map all gradients to their keys.
|
||||
named_grads = {
|
||||
module_id: {
|
||||
name: grad for (name, _), grad in zip(module_params.items(), grads)
|
||||
}
|
||||
for module_id, module_params in params.items()
|
||||
}
|
||||
|
||||
return named_grads
|
||||
|
||||
@override(DifferentiableLearner)
|
||||
def apply_gradients(
|
||||
self,
|
||||
gradients: Dict[ModuleID, NamedParamDict],
|
||||
params: Dict[ModuleID, NamedParamDict],
|
||||
) -> Dict[ModuleID, NamedParamDict]:
|
||||
"""Applies the given gradients in a functional manner.
|
||||
|
||||
This method requires functional parameter updates, meaning modifications
|
||||
must not be performed in-place (e.g., using an optimizer or directly within
|
||||
the `MultiRLModule`).
|
||||
|
||||
Args:
|
||||
gradients: A dictionary containing named gradients for each module id.
|
||||
params: A dictionary containing named parameters for each module id.
|
||||
|
||||
Returns:
|
||||
The updated parameters in the same (dict) format as `params`.
|
||||
"""
|
||||
policies_to_update = self.learner_config.policies_to_update or list(
|
||||
gradients.keys()
|
||||
)
|
||||
# Note, because this is a functional update we cannot apply in-place
|
||||
# modifications of parameters.
|
||||
updated_params = {}
|
||||
for module_id, module_grads in gradients.items():
|
||||
if module_id not in policies_to_update:
|
||||
updated_params[module_id] = params[module_id]
|
||||
continue
|
||||
updated_params[module_id] = {}
|
||||
for name, grad in module_grads.items():
|
||||
# If updates should not be skipped turn `nan` and `inf` gradients to zero.
|
||||
if (
|
||||
not self.config.torch_skip_nan_gradients
|
||||
and not torch.isfinite(grad).all()
|
||||
):
|
||||
# Warn the user about `nan` gradients.
|
||||
logger.warning(f"Gradients {name} contain `nan/inf` values.")
|
||||
# If updates should be skipped, do not step the optimizer and return.
|
||||
if not self.config.torch_skip_nan_gradients:
|
||||
logger.warning(
|
||||
"Setting `nan/inf` gradients to zero. If updates with "
|
||||
"`nan/inf` gradients should not be set to zero and instead "
|
||||
"the update be skipped entirely set `torch_skip_nan_gradients` "
|
||||
"to `True`."
|
||||
)
|
||||
# If necessary turn `nan` gradients to zero. Note, this can corrupt the
|
||||
# internal state of the optimizer, if many `nan` gradients occur.
|
||||
grad = torch.nan_to_num(grad)
|
||||
|
||||
if self.config.torch_skip_nan_gradients or torch.isfinite(grad).all():
|
||||
# Update each parameter, by a simple gradient descent step.
|
||||
updated_params[module_id][name] = (
|
||||
params[module_id][name] - self.learner_config.lr * grad
|
||||
)
|
||||
elif grad is None or not torch.isfinite(grad).all():
|
||||
logger.warning(
|
||||
"Skipping this update. If updates with `nan/inf` gradients "
|
||||
"should not be skipped entirely and instead `nan/inf` "
|
||||
"gradients set to `zero` set `torch_skip_nan_gradients` to "
|
||||
"`False`."
|
||||
)
|
||||
return updated_params
|
||||
|
||||
def _make_functional_call(
|
||||
self, params: Dict[ModuleID, NamedParamDict], batch: MultiAgentBatch
|
||||
) -> Dict[ModuleID, NamedParamDict]:
|
||||
"""Makes a functional call for each module in the `MultiRLModule`."""
|
||||
return self._module.foreach_module(
|
||||
lambda mid, m: torch.func.functional_call(m, params[mid], batch[mid]),
|
||||
return_dict=True,
|
||||
)
|
||||
|
||||
@override(DifferentiableLearner)
|
||||
def _get_tensor_variable(
|
||||
self, value, dtype=None, trainable=False
|
||||
) -> "torch.Tensor":
|
||||
tensor = torch.tensor(
|
||||
value,
|
||||
requires_grad=trainable,
|
||||
# TODO (simon): Make GPU-trainable.
|
||||
# device=self._device,
|
||||
dtype=(
|
||||
dtype
|
||||
or (
|
||||
torch.float32
|
||||
if isinstance(value, float)
|
||||
else torch.int32
|
||||
if isinstance(value, int)
|
||||
else None
|
||||
)
|
||||
),
|
||||
)
|
||||
return nn.Parameter(tensor) if trainable else tensor
|
||||
|
||||
def _convert_batch_type(
|
||||
self,
|
||||
batch: MultiAgentBatch,
|
||||
to_device: bool = True,
|
||||
pin_memory: bool = False,
|
||||
use_stream: bool = False,
|
||||
) -> MultiAgentBatch:
|
||||
batch = convert_to_torch_tensor(
|
||||
batch.policy_batches,
|
||||
device=self._device if to_device else None,
|
||||
pin_memory=pin_memory,
|
||||
use_stream=use_stream,
|
||||
)
|
||||
# TODO (sven): This computation of `env_steps` is not accurate!
|
||||
length = max(len(b) for b in batch.values())
|
||||
batch = MultiAgentBatch(batch, env_steps=length)
|
||||
return batch
|
||||
|
||||
def _compute_off_policyness(self, batch):
|
||||
# Log off-policy'ness of this batch wrt the current weights.
|
||||
off_policyness = {
|
||||
(mid, DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY): (
|
||||
(self._weights_seq_no - module_batch[WEIGHTS_SEQ_NO]).float()
|
||||
)
|
||||
for mid, module_batch in batch.items()
|
||||
if WEIGHTS_SEQ_NO in module_batch
|
||||
}
|
||||
for key in off_policyness.keys():
|
||||
mid = key[0]
|
||||
if Columns.LOSS_MASK not in batch[mid]:
|
||||
off_policyness[key] = torch.mean(off_policyness[key])
|
||||
else:
|
||||
mask = batch[mid][Columns.LOSS_MASK]
|
||||
num_valid = torch.sum(mask)
|
||||
off_policyness[key] = torch.sum(off_policyness[key][mask]) / num_valid
|
||||
self.metrics.log_dict(off_policyness, window=1)
|
||||
|
||||
@override(DifferentiableLearner)
|
||||
def build(self, device: Optional[DeviceType] = None) -> None:
|
||||
"""Builds the TorchDifferentiableLearner.
|
||||
|
||||
This method is specific to TorchDifferentiableLearner. Before running super() it will
|
||||
initialize the device properly based on `self.config`, so that `_make_module()`
|
||||
can place the created module on the correct device. After running super() it
|
||||
wraps the module in a TorchDDPRLModule if `config.num_learners > 0`.
|
||||
Note, in inherited classes it is advisable to call the parent's `build()`
|
||||
after setting up all variables because `configure_optimizer_for_module` is
|
||||
called in this `Learner.build()`.
|
||||
"""
|
||||
# TODO (simon): Allow different `DifferentiableLearner` instances in a
|
||||
# `MetaLearner` to run on different GPUs.
|
||||
|
||||
super().build(device=device)
|
||||
|
||||
if self._torch_compile_complete_update:
|
||||
torch._dynamo.reset()
|
||||
self._compiled_update_initialized = False
|
||||
self._possibly_compiled_update = torch.compile(
|
||||
self._uncompiled_update,
|
||||
backend=self._torch_compile_cfg.torch_dynamo_backend,
|
||||
mode=self._torch_compile_cfg.torch_dynamo_mode,
|
||||
**self._torch_compile_cfg.kwargs,
|
||||
)
|
||||
# Otherwise, we use the possibly compiled `forward_train` in
|
||||
# the module, compiled in the `MetaLearner`.
|
||||
else:
|
||||
# Nothing, to do.
|
||||
self._possibly_compiled_update = self._uncompiled_update
|
||||
|
||||
@override(DifferentiableLearner)
|
||||
def _update(
|
||||
self, batch: Dict[str, Any], params: Dict[ModuleID, NamedParamDict]
|
||||
) -> Tuple[Any, Dict[ModuleID, NamedParamDict], Any, Any]:
|
||||
# The first time we call _update after building the learner or
|
||||
# adding/removing models, we update with the uncompiled update method.
|
||||
# This makes it so that any variables that may be created during the first
|
||||
# update step are already there when compiling. More specifically,
|
||||
# this avoids errors that occur around using defaultdicts with
|
||||
# torch.compile().
|
||||
if (
|
||||
self._torch_compile_complete_update
|
||||
and not self._compiled_update_initialized
|
||||
):
|
||||
self._compiled_update_initialized = True
|
||||
return self._uncompiled_update(batch, params)
|
||||
else:
|
||||
return self._possibly_compiled_update(batch, params)
|
||||
@@ -0,0 +1,697 @@
|
||||
import contextlib
|
||||
import logging
|
||||
from collections import defaultdict
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Callable,
|
||||
Dict,
|
||||
Hashable,
|
||||
Optional,
|
||||
Sequence,
|
||||
Tuple,
|
||||
)
|
||||
|
||||
from ray.rllib.algorithms.algorithm_config import (
|
||||
AlgorithmConfig,
|
||||
TorchCompileWhatToCompile,
|
||||
)
|
||||
from ray.rllib.core.columns import Columns
|
||||
from ray.rllib.core.learner.learner import LR_KEY, Learner
|
||||
from ray.rllib.core.rl_module.multi_rl_module import (
|
||||
MultiRLModule,
|
||||
MultiRLModuleSpec,
|
||||
)
|
||||
from ray.rllib.core.rl_module.rl_module import (
|
||||
RLModule,
|
||||
RLModuleSpec,
|
||||
)
|
||||
from ray.rllib.core.rl_module.torch.torch_rl_module import (
|
||||
TorchCompileConfig,
|
||||
TorchDDPRLModule,
|
||||
TorchRLModule,
|
||||
)
|
||||
from ray.rllib.policy.sample_batch import MultiAgentBatch
|
||||
from ray.rllib.utils.annotations import (
|
||||
OverrideToImplementCustomLogic,
|
||||
OverrideToImplementCustomLogic_CallToSuperRecommended,
|
||||
override,
|
||||
)
|
||||
from ray.rllib.utils.framework import get_device, try_import_torch
|
||||
from ray.rllib.utils.metrics import (
|
||||
ALL_MODULES,
|
||||
DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY,
|
||||
NUM_NON_TRAINABLE_PARAMETERS,
|
||||
NUM_TRAINABLE_PARAMETERS,
|
||||
WEIGHTS_SEQ_NO,
|
||||
)
|
||||
from ray.rllib.utils.numpy import convert_to_numpy
|
||||
from ray.rllib.utils.torch_utils import convert_to_torch_tensor
|
||||
from ray.rllib.utils.typing import (
|
||||
ModuleID,
|
||||
Optimizer,
|
||||
Param,
|
||||
ParamDict,
|
||||
ShouldModuleBeUpdatedFn,
|
||||
StateDict,
|
||||
TensorType,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
|
||||
|
||||
torch, nn = try_import_torch()
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class TorchLearner(Learner):
|
||||
|
||||
framework: str = "torch"
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
# Whether to compile the RL Module of this learner. This implies that the.
|
||||
# forward_train method of the RL Module will be compiled. Further more,
|
||||
# other forward methods of the RL Module will be compiled on demand.
|
||||
# This is assumed to not happen, since other forwrad methods are not expected
|
||||
# to be used during training.
|
||||
self._torch_compile_forward_train = False
|
||||
self._torch_compile_cfg = None
|
||||
# Whether to compile the `_uncompiled_update` method of this learner. This
|
||||
# implies that everything within `_uncompiled_update` will be compiled,
|
||||
# not only the forward_train method of the RL Module.
|
||||
# Note that this is experimental.
|
||||
# Note that this requires recompiling the forward methods once we add/remove
|
||||
# RL Modules.
|
||||
self._torch_compile_complete_update = False
|
||||
if self.config.torch_compile_learner:
|
||||
if (
|
||||
self.config.torch_compile_learner_what_to_compile
|
||||
== TorchCompileWhatToCompile.COMPLETE_UPDATE
|
||||
):
|
||||
self._torch_compile_complete_update = True
|
||||
self._compiled_update_initialized = False
|
||||
else:
|
||||
self._torch_compile_forward_train = True
|
||||
|
||||
self._torch_compile_cfg = TorchCompileConfig(
|
||||
torch_dynamo_backend=self.config.torch_compile_learner_dynamo_backend,
|
||||
torch_dynamo_mode=self.config.torch_compile_learner_dynamo_mode,
|
||||
)
|
||||
|
||||
# Loss scalers for mixed precision training. Map optimizer names to
|
||||
# associated torch GradScaler objects.
|
||||
self._grad_scalers = None
|
||||
if self.config._torch_grad_scaler_class:
|
||||
self._grad_scalers = defaultdict(
|
||||
lambda: self.config._torch_grad_scaler_class()
|
||||
)
|
||||
self._lr_schedulers = {}
|
||||
self._lr_scheduler_classes = None
|
||||
if self.config._torch_lr_scheduler_classes:
|
||||
self._lr_scheduler_classes = self.config._torch_lr_scheduler_classes
|
||||
|
||||
@OverrideToImplementCustomLogic
|
||||
@override(Learner)
|
||||
def configure_optimizers_for_module(
|
||||
self,
|
||||
module_id: ModuleID,
|
||||
config: "AlgorithmConfig" = None,
|
||||
) -> None:
|
||||
module = self._module[module_id]
|
||||
|
||||
# For this default implementation, the learning rate is handled by the
|
||||
# attached lr Scheduler (controlled by self.config.lr, which can be a
|
||||
# fixed value or a schedule setting).
|
||||
params = self.get_parameters(module)
|
||||
optimizer = torch.optim.Adam(params)
|
||||
|
||||
# Register the created optimizer (under the default optimizer name).
|
||||
self.register_optimizer(
|
||||
module_id=module_id,
|
||||
optimizer=optimizer,
|
||||
params=params,
|
||||
lr_or_lr_schedule=config.lr,
|
||||
)
|
||||
|
||||
def _uncompiled_update(
|
||||
self,
|
||||
batch: Dict,
|
||||
**kwargs,
|
||||
):
|
||||
"""Performs a single update given a batch of data."""
|
||||
|
||||
# TODO (sven): Causes weird cuda error when WandB is used.
|
||||
# Diagnosis thus far:
|
||||
# - All peek values during metrics.reduce are non-tensors.
|
||||
# - However, in impala.py::training_step(), a tensor does arrive after learner
|
||||
# group.update(), so somehow, there is still a race condition
|
||||
# possible (learner, which performs the reduce() and learner thread, which
|
||||
# performs the logging of tensors into metrics logger).
|
||||
self._compute_off_policyness(batch)
|
||||
|
||||
fwd_out = self.module.forward_train(batch)
|
||||
loss_per_module = self.compute_losses(fwd_out=fwd_out, batch=batch)
|
||||
gradients = self.compute_gradients(loss_per_module)
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
if self.config.num_learners > 1:
|
||||
for mod in self.module.values():
|
||||
# Skip non-torch modules, b/c they may not have the `no_sync` API.
|
||||
if isinstance(mod, torch.nn.Module):
|
||||
stack.enter_context(mod.no_sync())
|
||||
postprocessed_gradients = self.postprocess_gradients(gradients)
|
||||
self.apply_gradients(postprocessed_gradients)
|
||||
|
||||
return fwd_out, loss_per_module, {}
|
||||
|
||||
@override(Learner)
|
||||
def compute_gradients(
|
||||
self, loss_per_module: Dict[ModuleID, TensorType], **kwargs
|
||||
) -> ParamDict:
|
||||
for optim in self._optimizer_parameters:
|
||||
# `set_to_none=True` is a faster way to zero out the gradients.
|
||||
optim.zero_grad(set_to_none=True)
|
||||
|
||||
if self._grad_scalers is not None:
|
||||
total_loss = sum(
|
||||
self._grad_scalers[mid].scale(loss)
|
||||
for mid, loss in loss_per_module.items()
|
||||
)
|
||||
else:
|
||||
total_loss = sum(loss_per_module.values())
|
||||
|
||||
# If we don't have any loss computations, `sum` returns 0.
|
||||
if isinstance(total_loss, int):
|
||||
assert total_loss == 0
|
||||
return {}
|
||||
|
||||
# If all parameters for this module are currently frozen, we can't optimize it.
|
||||
if total_loss.grad_fn is None:
|
||||
return {}
|
||||
|
||||
total_loss.backward()
|
||||
grads = {pid: p.grad for pid, p in self._params.items()}
|
||||
|
||||
return grads
|
||||
|
||||
@override(Learner)
|
||||
def apply_gradients(self, gradients_dict: ParamDict) -> None:
|
||||
# Set the gradient of the parameters.
|
||||
for pid, grad in gradients_dict.items():
|
||||
# If updates should not be skipped turn `nan` and `inf` gradients to zero.
|
||||
if (
|
||||
not self.config.torch_skip_nan_gradients
|
||||
and not torch.isfinite(grad).all()
|
||||
):
|
||||
# Warn the user about `nan` gradients.
|
||||
logger.warning(f"Gradients {pid} contain `nan/inf` values.")
|
||||
# If updates should be skipped, do not step the optimizer and return.
|
||||
if not self.config.torch_skip_nan_gradients:
|
||||
logger.warning(
|
||||
"Setting `nan/inf` gradients to zero. If updates with "
|
||||
"`nan/inf` gradients should not be set to zero and instead "
|
||||
"the update be skipped entirely set `torch_skip_nan_gradients` "
|
||||
"to `True`."
|
||||
)
|
||||
# If necessary turn `nan` gradients to zero. Note this can corrupt the
|
||||
# internal state of the optimizer, if many `nan` gradients occur.
|
||||
self._params[pid].grad = torch.nan_to_num(grad)
|
||||
# Otherwise, use the gradient as is.
|
||||
else:
|
||||
self._params[pid].grad = grad
|
||||
|
||||
# For each optimizer call its step function.
|
||||
for module_id, optimizer_names in self._module_optimizers.items():
|
||||
for optimizer_name in optimizer_names:
|
||||
optim = self.get_optimizer(module_id, optimizer_name)
|
||||
# If we have learning rate schedulers for a module add them, if
|
||||
# necessary.
|
||||
if self._lr_scheduler_classes is not None:
|
||||
if (
|
||||
module_id not in self._lr_schedulers
|
||||
or optimizer_name not in self._lr_schedulers[module_id]
|
||||
):
|
||||
# Set for each module and optimizer a scheduler.
|
||||
self._lr_schedulers[module_id] = {optimizer_name: []}
|
||||
# If the classes are in a dictionary each module might have
|
||||
# a different set of schedulers.
|
||||
if isinstance(self._lr_scheduler_classes, dict):
|
||||
scheduler_classes = self._lr_scheduler_classes[module_id]
|
||||
# Else, each module has the same learning rate schedulers.
|
||||
else:
|
||||
scheduler_classes = self._lr_scheduler_classes
|
||||
# Initialize and add the schedulers.
|
||||
for scheduler_class in scheduler_classes:
|
||||
self._lr_schedulers[module_id][optimizer_name].append(
|
||||
scheduler_class(optim)
|
||||
)
|
||||
|
||||
# Step through the scaler (unscales gradients, if applicable).
|
||||
if self._grad_scalers is not None:
|
||||
scaler = self._grad_scalers[module_id]
|
||||
scaler.step(optim)
|
||||
self.metrics.log_value(
|
||||
(module_id, "_torch_grad_scaler_current_scale"),
|
||||
scaler.get_scale(),
|
||||
window=1, # snapshot in time, no EMA/mean.
|
||||
)
|
||||
# Update the scaler.
|
||||
scaler.update()
|
||||
# `step` the optimizer (default), but only if all gradients are finite.
|
||||
elif self.config.torch_skip_nan_gradients or all(
|
||||
param.grad is None or torch.isfinite(param.grad).all()
|
||||
for group in optim.param_groups
|
||||
for param in group["params"]
|
||||
):
|
||||
optim.step()
|
||||
# If gradients are not all finite warn the user that the update will be
|
||||
# skipped.
|
||||
elif not all(
|
||||
torch.isfinite(param.grad).all()
|
||||
for group in optim.param_groups
|
||||
for param in group["params"]
|
||||
):
|
||||
logger.warning(
|
||||
"Skipping this update. If updates with `nan/inf` gradients "
|
||||
"should not be skipped entirely and instead `nan/inf` "
|
||||
"gradients set to `zero` set `torch_skip_nan_gradients` to "
|
||||
"`False`."
|
||||
)
|
||||
|
||||
@OverrideToImplementCustomLogic_CallToSuperRecommended
|
||||
@override(Learner)
|
||||
def after_gradient_based_update(self, *, timesteps: Dict[str, Any]) -> None:
|
||||
"""Called after gradient-based updates are completed.
|
||||
|
||||
Should be overridden to implement custom cleanup-, logging-, or non-gradient-
|
||||
based Learner/RLModule update logic after(!) gradient-based updates have been
|
||||
completed.
|
||||
|
||||
Note, for `framework="torch"` users can register
|
||||
`torch.optim.lr_scheduler.LRScheduler` via
|
||||
`AlgorithmConfig._torch_lr_scheduler_classes`. These schedulers need to be
|
||||
stepped here after gradient updates and reported.
|
||||
|
||||
Args:
|
||||
timesteps: Timesteps dict, which must have the key
|
||||
`NUM_ENV_STEPS_SAMPLED_LIFETIME`.
|
||||
# TODO (sven): Make this a more formal structure with its own type.
|
||||
"""
|
||||
|
||||
# If we have no `torch.optim.lr_scheduler.LRScheduler` registered call the
|
||||
# `super()`'s method to update RLlib's learning rate schedules.
|
||||
if not self._lr_schedulers:
|
||||
return super().after_gradient_based_update(timesteps=timesteps)
|
||||
|
||||
# Only update this optimizer's lr, if a scheduler has been registered
|
||||
# along with it.
|
||||
for module_id, optimizer_names in self._module_optimizers.items():
|
||||
for optimizer_name in optimizer_names:
|
||||
# If learning rate schedulers are provided step them here. Note,
|
||||
# stepping them in `TorchLearner.apply_gradients` updates the
|
||||
# learning rates during minibatch updates; we want to update
|
||||
# between whole batch updates.
|
||||
if (
|
||||
module_id in self._lr_schedulers
|
||||
and optimizer_name in self._lr_schedulers[module_id]
|
||||
):
|
||||
for scheduler in self._lr_schedulers[module_id][optimizer_name]:
|
||||
scheduler.step()
|
||||
optimizer = self.get_optimizer(module_id, optimizer_name)
|
||||
self.metrics.log_value(
|
||||
# Cut out the module ID from the beginning since it's already
|
||||
# part of the key sequence: (ModuleID, "[optim name]_lr").
|
||||
key=(
|
||||
module_id,
|
||||
f"{optimizer_name[len(module_id) + 1:]}_{LR_KEY}",
|
||||
),
|
||||
value=self._get_optimizer_lr(optimizer),
|
||||
window=1,
|
||||
)
|
||||
|
||||
@override(Learner)
|
||||
def _get_optimizer_state(self) -> StateDict:
|
||||
ret = {}
|
||||
for name, optim in self._named_optimizers.items():
|
||||
ret[name] = {
|
||||
"module_id": self._optimizer_name_to_module[name],
|
||||
"state": convert_to_numpy(optim.state_dict()),
|
||||
}
|
||||
return ret
|
||||
|
||||
@override(Learner)
|
||||
def _set_optimizer_state(self, state: StateDict) -> None:
|
||||
for name, state_dict in state.items():
|
||||
# Ignore updating optimizers matching to submodules not present in this
|
||||
# Learner's MultiRLModule.
|
||||
module_id = state_dict["module_id"]
|
||||
if name not in self._named_optimizers and module_id in self.module:
|
||||
self.configure_optimizers_for_module(
|
||||
module_id=module_id,
|
||||
config=self.config.get_config_for_module(module_id=module_id),
|
||||
)
|
||||
if name in self._named_optimizers:
|
||||
# Keep optimizer param_groups untouched so scalar metadata such as
|
||||
# betas, lr, eps, foreach, capturable preserve their Python types.
|
||||
# Only convert per-parameter optimizer buffers (exp_avg, exp_avg_sq,
|
||||
# step, etc.) to tensors on the learner device.
|
||||
optimizer_state = state_dict["state"]
|
||||
loaded_state = {
|
||||
**optimizer_state,
|
||||
"state": convert_to_torch_tensor(
|
||||
optimizer_state["state"],
|
||||
device=self._device,
|
||||
),
|
||||
}
|
||||
self._named_optimizers[name].load_state_dict(loaded_state)
|
||||
|
||||
@override(Learner)
|
||||
def get_param_ref(self, param: Param) -> Hashable:
|
||||
return param
|
||||
|
||||
@override(Learner)
|
||||
def get_parameters(self, module: RLModule) -> Sequence[Param]:
|
||||
return list(module.parameters())
|
||||
|
||||
@override(Learner)
|
||||
def _convert_batch_type(
|
||||
self,
|
||||
batch: MultiAgentBatch,
|
||||
to_device: bool = True,
|
||||
pin_memory: bool = False,
|
||||
use_stream: bool = False,
|
||||
) -> MultiAgentBatch:
|
||||
batch_dict = convert_to_torch_tensor(
|
||||
batch.policy_batches,
|
||||
device=self._device if to_device else None,
|
||||
pin_memory=pin_memory,
|
||||
use_stream=use_stream,
|
||||
)
|
||||
batch = MultiAgentBatch(batch_dict, env_steps=batch.env_steps())
|
||||
return batch
|
||||
|
||||
@override(Learner)
|
||||
def add_module(
|
||||
self,
|
||||
*,
|
||||
module_id: ModuleID,
|
||||
# TODO (sven): Rename to `rl_module_spec`.
|
||||
module_spec: RLModuleSpec,
|
||||
config_overrides: Optional[Dict] = None,
|
||||
new_should_module_be_updated: Optional[ShouldModuleBeUpdatedFn] = None,
|
||||
) -> MultiRLModuleSpec:
|
||||
# Call super's add_module method.
|
||||
marl_spec = 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,
|
||||
)
|
||||
|
||||
# we need to ddpify the module that was just added to the pool
|
||||
module = self._module[module_id]
|
||||
|
||||
if self._torch_compile_forward_train:
|
||||
module.compile(self._torch_compile_cfg)
|
||||
elif self._torch_compile_complete_update:
|
||||
# When compiling the update, we need to reset and recompile
|
||||
# _uncompiled_update every time we add/remove a module anew.
|
||||
torch._dynamo.reset()
|
||||
self._compiled_update_initialized = False
|
||||
self._possibly_compiled_update = torch.compile(
|
||||
self._uncompiled_update,
|
||||
backend=self._torch_compile_cfg.torch_dynamo_backend,
|
||||
mode=self._torch_compile_cfg.torch_dynamo_mode,
|
||||
**self._torch_compile_cfg.kwargs,
|
||||
)
|
||||
|
||||
if isinstance(module, TorchRLModule):
|
||||
self._module[module_id].to(self._device)
|
||||
if self.distributed:
|
||||
if (
|
||||
self._torch_compile_complete_update
|
||||
or self._torch_compile_forward_train
|
||||
):
|
||||
raise ValueError(
|
||||
"Using torch distributed and torch compile "
|
||||
"together tested for now. Please disable "
|
||||
"torch compile."
|
||||
)
|
||||
self._module.add_module(
|
||||
module_id,
|
||||
TorchDDPRLModule(module, **self.config.torch_ddp_kwargs),
|
||||
override=True,
|
||||
)
|
||||
|
||||
self._log_trainable_parameters()
|
||||
|
||||
return marl_spec
|
||||
|
||||
@override(Learner)
|
||||
def remove_module(self, module_id: ModuleID, **kwargs) -> MultiRLModuleSpec:
|
||||
marl_spec = super().remove_module(module_id, **kwargs)
|
||||
|
||||
if self._torch_compile_complete_update:
|
||||
# When compiling the update, we need to reset and recompile
|
||||
# _uncompiled_update every time we add/remove a module anew.
|
||||
torch._dynamo.reset()
|
||||
self._compiled_update_initialized = False
|
||||
self._possibly_compiled_update = torch.compile(
|
||||
self._uncompiled_update,
|
||||
backend=self._torch_compile_cfg.torch_dynamo_backend,
|
||||
mode=self._torch_compile_cfg.torch_dynamo_mode,
|
||||
**self._torch_compile_cfg.kwargs,
|
||||
)
|
||||
|
||||
self._log_trainable_parameters()
|
||||
|
||||
return marl_spec
|
||||
|
||||
@override(Learner)
|
||||
def build(self) -> None:
|
||||
"""Builds the TorchLearner.
|
||||
|
||||
This method is specific to TorchLearner. Before running super() it will
|
||||
initialize the device properly based on `self.config`, so that `_make_module()`
|
||||
can place the created module on the correct device. After running super() it
|
||||
wraps the module in a TorchDDPRLModule if `config.num_learners > 0`.
|
||||
Note, in inherited classes it is advisable to call the parent's `build()`
|
||||
after setting up all variables because `configure_optimizer_for_module` is
|
||||
called in this `Learner.build()`.
|
||||
"""
|
||||
self._device = get_device(self.config, self.config.num_gpus_per_learner)
|
||||
|
||||
super().build()
|
||||
|
||||
if self._torch_compile_complete_update:
|
||||
torch._dynamo.reset()
|
||||
self._compiled_update_initialized = False
|
||||
self._possibly_compiled_update = torch.compile(
|
||||
self._uncompiled_update,
|
||||
backend=self._torch_compile_cfg.torch_dynamo_backend,
|
||||
mode=self._torch_compile_cfg.torch_dynamo_mode,
|
||||
**self._torch_compile_cfg.kwargs,
|
||||
)
|
||||
else:
|
||||
if self._torch_compile_forward_train:
|
||||
if isinstance(self._module, TorchRLModule):
|
||||
self._module.compile(self._torch_compile_cfg)
|
||||
elif isinstance(self._module, MultiRLModule):
|
||||
for module in self._module._rl_modules.values():
|
||||
# Compile only TorchRLModules, e.g. we don't want to compile
|
||||
# a RandomRLModule.
|
||||
if isinstance(self._module, TorchRLModule):
|
||||
module.compile(self._torch_compile_cfg)
|
||||
else:
|
||||
raise ValueError(
|
||||
"Torch compile is only supported for TorchRLModule and "
|
||||
"MultiRLModule."
|
||||
)
|
||||
|
||||
self._possibly_compiled_update = self._uncompiled_update
|
||||
|
||||
self._make_modules_ddp_if_necessary()
|
||||
|
||||
@override(Learner)
|
||||
def _update(self, batch: Dict[str, Any]) -> Tuple[Any, Any, Any]:
|
||||
# The first time we call _update after building the learner or
|
||||
# adding/removing models, we update with the uncompiled update method.
|
||||
# This makes it so that any variables that may be created during the first
|
||||
# update step are already there when compiling. More specifically,
|
||||
# this avoids errors that occur around using defaultdicts with
|
||||
# torch.compile().
|
||||
if (
|
||||
self._torch_compile_complete_update
|
||||
and not self._compiled_update_initialized
|
||||
):
|
||||
self._compiled_update_initialized = True
|
||||
return self._uncompiled_update(batch)
|
||||
else:
|
||||
return self._possibly_compiled_update(batch)
|
||||
|
||||
@OverrideToImplementCustomLogic
|
||||
def _make_modules_ddp_if_necessary(self) -> None:
|
||||
"""Default logic for (maybe) making all Modules within self._module DDP."""
|
||||
|
||||
# If the module is a MultiRLModule and nn.Module we can simply assume
|
||||
# all the submodules are registered. Otherwise, we need to loop through
|
||||
# each submodule and move it to the correct device.
|
||||
# TODO (Kourosh): This can result in missing modules if the user does not
|
||||
# register them in the MultiRLModule. We should find a better way to
|
||||
# handle this.
|
||||
if self.config.num_learners > 1:
|
||||
# Single agent module: Convert to `TorchDDPRLModule`.
|
||||
if isinstance(self._module, TorchRLModule):
|
||||
self._module = TorchDDPRLModule(
|
||||
self._module, **self.config.torch_ddp_kwargs
|
||||
)
|
||||
# Multi agent module: Convert each submodule to `TorchDDPRLModule`.
|
||||
else:
|
||||
assert isinstance(self._module, MultiRLModule)
|
||||
for key in self._module.keys():
|
||||
sub_module = self._module[key]
|
||||
if isinstance(sub_module, TorchRLModule):
|
||||
# Wrap and override the module ID key in self._module.
|
||||
self._module.add_module(
|
||||
key,
|
||||
TorchDDPRLModule(
|
||||
sub_module, **self.config.torch_ddp_kwargs
|
||||
),
|
||||
override=True,
|
||||
)
|
||||
|
||||
def rl_module_is_compatible(self, module: RLModule) -> bool:
|
||||
return isinstance(module, nn.Module)
|
||||
|
||||
@override(Learner)
|
||||
def _check_registered_optimizer(
|
||||
self,
|
||||
optimizer: Optimizer,
|
||||
params: Sequence[Param],
|
||||
) -> None:
|
||||
super()._check_registered_optimizer(optimizer, params)
|
||||
if not isinstance(optimizer, torch.optim.Optimizer):
|
||||
raise ValueError(
|
||||
f"The optimizer ({optimizer}) is not a torch.optim.Optimizer! "
|
||||
"Only use torch.optim.Optimizer subclasses for TorchLearner."
|
||||
)
|
||||
for param in params:
|
||||
if not isinstance(param, torch.Tensor):
|
||||
raise ValueError(
|
||||
f"One of the parameters ({param}) in the registered optimizer "
|
||||
"is not a torch.Tensor!"
|
||||
)
|
||||
|
||||
@override(Learner)
|
||||
def _make_module(self) -> MultiRLModule:
|
||||
module = super()._make_module()
|
||||
self._map_module_to_device(module)
|
||||
return module
|
||||
|
||||
def _map_module_to_device(self, module: MultiRLModule) -> None:
|
||||
"""Moves the module to the correct device."""
|
||||
if isinstance(module, torch.nn.Module):
|
||||
module.to(self._device)
|
||||
else:
|
||||
for key in module.keys():
|
||||
if isinstance(module[key], torch.nn.Module):
|
||||
module[key].to(self._device)
|
||||
|
||||
@override(Learner)
|
||||
def _log_trainable_parameters(self) -> None:
|
||||
# Log number of non-trainable and trainable parameters of our RLModule.
|
||||
num_trainable_params = defaultdict(int)
|
||||
num_non_trainable_params = defaultdict(int)
|
||||
for mid, rlm in self.module._rl_modules.items():
|
||||
if isinstance(rlm, TorchRLModule):
|
||||
for p in rlm.parameters():
|
||||
n = p.numel()
|
||||
if p.requires_grad:
|
||||
num_trainable_params[(mid, NUM_TRAINABLE_PARAMETERS)] += n
|
||||
else:
|
||||
num_non_trainable_params[
|
||||
(mid, NUM_NON_TRAINABLE_PARAMETERS)
|
||||
] += n
|
||||
|
||||
self.metrics.log_dict(
|
||||
{
|
||||
**{
|
||||
(ALL_MODULES, NUM_TRAINABLE_PARAMETERS): sum(
|
||||
num_trainable_params.values()
|
||||
),
|
||||
(ALL_MODULES, NUM_NON_TRAINABLE_PARAMETERS): sum(
|
||||
num_non_trainable_params.values()
|
||||
),
|
||||
},
|
||||
**num_trainable_params,
|
||||
**num_non_trainable_params,
|
||||
}
|
||||
)
|
||||
|
||||
def _compute_off_policyness(self, batch):
|
||||
# Log off-policy'ness of this batch wrt the current weights.
|
||||
off_policyness = {
|
||||
(mid, DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY): (
|
||||
(self._weights_seq_no - module_batch[WEIGHTS_SEQ_NO]).float()
|
||||
)
|
||||
for mid, module_batch in batch.items()
|
||||
if WEIGHTS_SEQ_NO in module_batch
|
||||
}
|
||||
for key in off_policyness.keys():
|
||||
mid = key[0]
|
||||
if Columns.LOSS_MASK not in batch[mid]:
|
||||
off_policyness[key] = torch.mean(off_policyness[key])
|
||||
else:
|
||||
mask = batch[mid][Columns.LOSS_MASK]
|
||||
num_valid = torch.sum(mask)
|
||||
off_policyness[key] = torch.sum(off_policyness[key][mask]) / num_valid
|
||||
self.metrics.log_dict(off_policyness, window=1)
|
||||
|
||||
@override(Learner)
|
||||
def _get_tensor_variable(
|
||||
self, value, dtype=None, trainable=False
|
||||
) -> "torch.Tensor":
|
||||
tensor = torch.tensor(
|
||||
value,
|
||||
requires_grad=trainable,
|
||||
device=self._device,
|
||||
dtype=(
|
||||
dtype
|
||||
or (
|
||||
torch.float32
|
||||
if isinstance(value, float)
|
||||
else torch.int32
|
||||
if isinstance(value, int)
|
||||
else None
|
||||
)
|
||||
),
|
||||
)
|
||||
return nn.Parameter(tensor) if trainable else tensor
|
||||
|
||||
@staticmethod
|
||||
@override(Learner)
|
||||
def _get_optimizer_lr(optimizer: "torch.optim.Optimizer") -> float:
|
||||
for g in optimizer.param_groups:
|
||||
return g["lr"]
|
||||
|
||||
@staticmethod
|
||||
@override(Learner)
|
||||
def _set_optimizer_lr(optimizer: "torch.optim.Optimizer", lr: float) -> None:
|
||||
for g in optimizer.param_groups:
|
||||
g["lr"] = lr
|
||||
|
||||
@staticmethod
|
||||
@override(Learner)
|
||||
def _get_clip_function() -> Callable:
|
||||
from ray.rllib.utils.torch_utils import clip_gradients
|
||||
|
||||
return clip_gradients
|
||||
|
||||
@staticmethod
|
||||
@override(Learner)
|
||||
def _get_global_norm_function() -> Callable:
|
||||
from ray.rllib.utils.torch_utils import compute_global_norm
|
||||
|
||||
return compute_global_norm
|
||||
@@ -0,0 +1,449 @@
|
||||
import contextlib
|
||||
import logging
|
||||
from itertools import cycle
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
import ray
|
||||
from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
|
||||
from ray.rllib.core import ALL_MODULES
|
||||
from ray.rllib.core.learner.learner import Learner
|
||||
from ray.rllib.core.learner.torch.torch_differentiable_learner import (
|
||||
TorchDifferentiableLearner,
|
||||
)
|
||||
from ray.rllib.core.learner.torch.torch_learner import TorchLearner
|
||||
from ray.rllib.core.learner.training_data import TrainingData
|
||||
from ray.rllib.core.rl_module.apis import SelfSupervisedLossAPI
|
||||
from ray.rllib.policy.sample_batch import MultiAgentBatch
|
||||
from ray.rllib.utils.annotations import (
|
||||
OverrideToImplementCustomLogic,
|
||||
OverrideToImplementCustomLogic_CallToSuperRecommended,
|
||||
override,
|
||||
)
|
||||
from ray.rllib.utils.framework import try_import_torch
|
||||
from ray.rllib.utils.metrics import (
|
||||
DATASET_NUM_ITERS_TRAINED,
|
||||
DATASET_NUM_ITERS_TRAINED_LIFETIME,
|
||||
DIFFERENTIABLE_LEARNER_RESULTS,
|
||||
)
|
||||
from ray.rllib.utils.metrics.utils import to_snake_case
|
||||
from ray.rllib.utils.typing import (
|
||||
EpisodeType,
|
||||
ModuleID,
|
||||
NamedParamDict,
|
||||
ResultDict,
|
||||
TensorType,
|
||||
)
|
||||
|
||||
logger = logging.getLogger("__name__")
|
||||
|
||||
torch, nn = try_import_torch()
|
||||
|
||||
|
||||
class TorchMetaLearner(TorchLearner):
|
||||
"""A `TorchLearner` designed for meta-learning with functional updates.
|
||||
|
||||
This `TorchLearner` manages one or more `DifferentiableLearner` instances,
|
||||
which perform functional updates on the `MultiRLModule`. These updates enable
|
||||
the computation of higher-order gradients, making this class suitable for
|
||||
meta-learning applications.
|
||||
|
||||
The `update` method executes one or more update loops on its
|
||||
`DifferentiableLearner` instances, leveraging the functionally updated
|
||||
parameters for its own learning process.
|
||||
"""
|
||||
|
||||
# A list of `TorchDifferentiableLearner`s that run inner functional update
|
||||
# loops.
|
||||
others: List[TorchDifferentiableLearner]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
**kwargs,
|
||||
):
|
||||
# First, initialize the `TorchLearner`.
|
||||
super().__init__(**kwargs)
|
||||
|
||||
# Initialize all configured `TorchDifferentiableLearner`s.
|
||||
self.others = [
|
||||
other_config.learner_class(
|
||||
config=self.config,
|
||||
learner_config=other_config,
|
||||
module=self.module,
|
||||
)
|
||||
for other_config in self.config.differentiable_learner_configs
|
||||
]
|
||||
|
||||
@OverrideToImplementCustomLogic_CallToSuperRecommended
|
||||
@override(TorchLearner)
|
||||
def build(self) -> None:
|
||||
"""Builds the `TorchMetaLearner`.
|
||||
|
||||
This method should be called before the meta-learner is used. It is
|
||||
responsible for setting up the `TorchDifferentiableLearner`s.
|
||||
"""
|
||||
super().build()
|
||||
|
||||
# Build all `DifferentiableLearner`s.
|
||||
for other in self.others:
|
||||
other.build(device=self._device)
|
||||
|
||||
# Ensure 'others' have a module reference.
|
||||
for other in self.others:
|
||||
if not other.module:
|
||||
other.module = self.module
|
||||
|
||||
@OverrideToImplementCustomLogic
|
||||
@override(Learner)
|
||||
def update(
|
||||
self,
|
||||
batch: Optional[MultiAgentBatch] = None,
|
||||
batches: Optional[List[MultiAgentBatch]] = None,
|
||||
batch_refs: Optional[List[ray.ObjectRef]] = None,
|
||||
episodes: Optional[List[EpisodeType]] = None,
|
||||
episodes_refs: Optional[List[ray.ObjectRef]] = None,
|
||||
data_iterators: Optional[List[ray.data.DataIterator]] = None,
|
||||
training_data: Optional[TrainingData] = None,
|
||||
*,
|
||||
# TODO (simon): Maybe defined here other_training_data to provide
|
||||
# TrainingData for differentiable Learners. List to nest them.
|
||||
# other_training_data: List[Union[List[Any], TrainingData]]
|
||||
# TODO (sven): Make this a more formal structure with its own type.
|
||||
timesteps: Optional[Dict[str, Any]] = None,
|
||||
num_total_minibatches: int = 0,
|
||||
num_epochs: int = 1,
|
||||
minibatch_size: Optional[int] = None,
|
||||
shuffle_batch_per_epoch: bool = False,
|
||||
_no_metrics_reduce: bool = False,
|
||||
others_training_data: Optional[List[TrainingData]] = None,
|
||||
**kwargs,
|
||||
) -> ResultDict:
|
||||
"""Performs a meta-update on the `MultiRLModule`.
|
||||
|
||||
This method allows multiple backward passes on a batch by iteratively
|
||||
applying functional updates from the `TorchDifferentiableLearner` instances
|
||||
in `others`. The resulting differentiable parameters are then used in this
|
||||
learner's own functional forward pass on the `MultiRLModule`, enabling the
|
||||
computation of higher-order gradients during the backward pass.
|
||||
"""
|
||||
self._check_is_built()
|
||||
|
||||
# Call `before_gradient_based_update` to allow for non-gradient based
|
||||
# preparations-, logging-, and update logic to happen.
|
||||
self.before_gradient_based_update(timesteps=timesteps or {})
|
||||
|
||||
if training_data is None:
|
||||
training_data = TrainingData(
|
||||
batch=batch,
|
||||
batches=batches,
|
||||
batch_refs=batch_refs,
|
||||
episodes=episodes,
|
||||
episodes_refs=episodes_refs,
|
||||
data_iterators=data_iterators,
|
||||
)
|
||||
# Validate training data.
|
||||
# TODO (simon): Pass in a `TrainingData` list for all differentiable learners.
|
||||
# to be used if necessary. Some algorithms need it.
|
||||
training_data.validate()
|
||||
training_data.solve_refs()
|
||||
assert training_data.batches is None, "`training_data.batches` must be None!"
|
||||
|
||||
# Increase the _weigths_seq_no in each `update` run.
|
||||
self._weights_seq_no += 1
|
||||
|
||||
# Create a batch iterator.
|
||||
# TODO (simon): Create this method in the `Learner`.
|
||||
batch_iter = self._create_iterator_if_necessary(
|
||||
training_data=training_data,
|
||||
num_total_minibatches=num_total_minibatches,
|
||||
num_epochs=num_epochs,
|
||||
minibatch_size=minibatch_size,
|
||||
shuffle_batch_per_epoch=shuffle_batch_per_epoch,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# If no training data for `DifferentiableLearner`s have been passed in, cycle
|
||||
# over the main training_data for each `DifferentiableLearner`.
|
||||
if not others_training_data:
|
||||
others_training_data = cycle([training_data])
|
||||
|
||||
# Perform the actual looping through the minibatches or the given data iterator.
|
||||
for iteration, tensor_minibatch in enumerate(batch_iter):
|
||||
# Check the MultiAgentBatch, whether our RLModule contains all ModuleIDs
|
||||
# found in this batch. If not, throw an error.
|
||||
unknown_module_ids = set(tensor_minibatch.policy_batches.keys()) - set(
|
||||
self.module.keys()
|
||||
)
|
||||
if unknown_module_ids:
|
||||
raise ValueError(
|
||||
f"Batch contains one or more ModuleIDs ({unknown_module_ids}) that "
|
||||
f"are not in this Learner!"
|
||||
)
|
||||
|
||||
# Clone the parameters for the differentiable updates.
|
||||
params = self._clone_named_parameters()
|
||||
|
||||
# Update all differentiable learners.
|
||||
others_loss_per_module = []
|
||||
others_results = {}
|
||||
for other, other_training_data in zip(self.others, others_training_data):
|
||||
# TODO (simon): Maybe return all losses from each other's update step.
|
||||
params, other_loss_per_module, other_results = other.update(
|
||||
training_data=other_training_data,
|
||||
params=params,
|
||||
# TODO (simon): Check, if this is still needed.
|
||||
_no_metrics_reduce=_no_metrics_reduce,
|
||||
**kwargs,
|
||||
)
|
||||
others_loss_per_module.append(other_loss_per_module)
|
||||
# TODO (simon): Find a more elegant way for naming.
|
||||
others_results[to_snake_case(other.__class__.__name__)] = other_results
|
||||
|
||||
# Log training results from the `DifferentiableLearner`s.
|
||||
# TODO (simon): Right now metrics are not carried over b/c of
|
||||
# the double tensormode problem.
|
||||
self.metrics.aggregate(
|
||||
stats_dicts=[others_results], key=DIFFERENTIABLE_LEARNER_RESULTS
|
||||
)
|
||||
|
||||
# Make the actual in-graph/traced meta-`_update` call. This should return
|
||||
# all tensor values (no numpy).
|
||||
fwd_out, loss_per_module, _ = self._update(
|
||||
tensor_minibatch.policy_batches,
|
||||
params,
|
||||
others_loss_per_module,
|
||||
)
|
||||
|
||||
# TODO (sven): Maybe move this into loop above to get metrics more accuratcely
|
||||
# cover the minibatch/epoch logic.
|
||||
# Log all timesteps (env, agent, modules) based on given episodes/batch.
|
||||
self._log_steps_trained_metrics(tensor_minibatch)
|
||||
|
||||
self._set_slicing_by_batch_id(tensor_minibatch, value=False)
|
||||
|
||||
if self.iterator:
|
||||
# Record the number of batches pulled from the dataset.
|
||||
self.metrics.log_value(
|
||||
(ALL_MODULES, DATASET_NUM_ITERS_TRAINED),
|
||||
iteration + 1,
|
||||
reduce="sum",
|
||||
)
|
||||
self.metrics.log_value(
|
||||
(ALL_MODULES, DATASET_NUM_ITERS_TRAINED_LIFETIME),
|
||||
iteration + 1,
|
||||
reduce="lifetime_sum",
|
||||
)
|
||||
# Log all individual RLModules' loss terms and its registered optimizers'
|
||||
# current learning rates.
|
||||
# Note: We do this only once for the last of the minibatch updates, b/c the
|
||||
# window is only 1 anyways.
|
||||
for mid, loss in loss_per_module.items():
|
||||
self.metrics.log_value(
|
||||
key=(mid, self.TOTAL_LOSS_KEY),
|
||||
value=loss,
|
||||
window=1,
|
||||
)
|
||||
|
||||
# Call `after_gradient_based_update` to allow for non-gradient based
|
||||
# cleanups-, logging-, and update logic to happen.
|
||||
# TODO (simon): Check, if this should stay here, when running multiple
|
||||
# gradient steps inside the iterator loop above (could be a complete epoch)
|
||||
# the target networks might need to be updated earlier.
|
||||
self.after_gradient_based_update(timesteps=timesteps or {})
|
||||
|
||||
# Reduce results across all minibatch update steps.
|
||||
if not _no_metrics_reduce:
|
||||
return self.metrics.reduce()
|
||||
|
||||
def _clone_named_parameters(self):
|
||||
"""Clone named parameters for functional updates."""
|
||||
return self.module.foreach_module(
|
||||
lambda _, m: {name: p.clone() for name, p in m.named_parameters()},
|
||||
return_dict=True,
|
||||
)
|
||||
|
||||
@OverrideToImplementCustomLogic
|
||||
@override(TorchLearner)
|
||||
def _update(
|
||||
self,
|
||||
batch: Dict[str, Any],
|
||||
params: NamedParamDict,
|
||||
others_loss_per_module: List[Dict[ModuleID, TensorType]] = None,
|
||||
) -> Tuple[Any, Any, Any]:
|
||||
# The first time we call _update after building the learner or
|
||||
# adding/removing models, we update with the uncompiled update method.
|
||||
# This makes it so that any variables that may be created during the first
|
||||
# update step are already there when compiling. More specifically,
|
||||
# this avoids errors that occur around using defaultdicts with
|
||||
# torch.compile().
|
||||
if (
|
||||
self._torch_compile_complete_update
|
||||
and not self._compiled_update_initialized
|
||||
):
|
||||
self._compiled_update_initialized = True
|
||||
return self._uncompiled_update(batch, params, others_loss_per_module)
|
||||
else:
|
||||
return self._possibly_compiled_update(batch, params, others_loss_per_module)
|
||||
|
||||
def _uncompiled_update(
|
||||
self,
|
||||
batch: Dict,
|
||||
params: NamedParamDict,
|
||||
others_loss_per_module: List[Dict[ModuleID, TensorType]] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""Performs a single functional update using a batch of data.
|
||||
|
||||
This update utilizes a functional forward pass via PyTorch 2.0's `func` module.
|
||||
The updated parameters provided by `TorchDifferentiableLearner` instances are
|
||||
functions of the `MultiRLModule`'s parameters, allowing for differentiation.
|
||||
|
||||
To ensure compatibility, the `MultiRLModule`'s `forward` method must encapsulate
|
||||
all logic from `forward_train` and support additional keyword arguments (`kwargs`).
|
||||
"""
|
||||
# TODO (sven): Causes weird cuda error when WandB is used.
|
||||
# Diagnosis thus far:
|
||||
# - All peek values during metrics.reduce are non-tensors.
|
||||
# - However, in impala.py::training_step(), a tensor does arrive after learner
|
||||
# group.update(), so somehow, there is still a race condition
|
||||
# possible (learner, which performs the reduce() and learner thread, which
|
||||
# performs the logging of tensors into metrics logger).
|
||||
self._compute_off_policyness(batch)
|
||||
|
||||
# TODO (simon): For this to work, the `RLModule.forward` must run the
|
||||
# the `forward_train`. Passing in arguments which makes the `forward` modular
|
||||
# could be a workaround.
|
||||
# Make a functional forward call to include higher-order gradients in the meta
|
||||
# update.
|
||||
fwd_out = self._make_functional_call(params, batch)
|
||||
loss_per_module = self.compute_losses(
|
||||
fwd_out=fwd_out, batch=batch, others_loss_per_module=others_loss_per_module
|
||||
)
|
||||
gradients = self.compute_gradients(loss_per_module)
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
if self.config.num_learners > 1:
|
||||
for mod in self.module.values():
|
||||
# Skip non-torch modules, b/c they may not have the `no_sync` API.
|
||||
if isinstance(mod, torch.nn.Module):
|
||||
stack.enter_context(mod.no_sync())
|
||||
postprocessed_gradients = self.postprocess_gradients(gradients)
|
||||
self.apply_gradients(postprocessed_gradients)
|
||||
|
||||
# Deactivate tensor-mode on our MetricsLogger and collect the (tensor)
|
||||
# results.
|
||||
return fwd_out, loss_per_module, {}
|
||||
|
||||
@override(Learner)
|
||||
def compute_losses(
|
||||
self,
|
||||
*,
|
||||
fwd_out: Dict[str, Any],
|
||||
batch: Dict[str, Any],
|
||||
others_loss_per_module: List[Dict[ModuleID, TensorType]] = None,
|
||||
**kwargs,
|
||||
) -> Dict[ModuleID, TensorType]:
|
||||
"""Computes the loss(es) for the module being optimized.
|
||||
|
||||
This method must be overridden by MultiRLModule-specific Learners in order to
|
||||
define the specific loss computation logic. If the algorithm is single-agent,
|
||||
only `compute_loss_for_module()` should be overridden instead. If the algorithm
|
||||
uses independent multi-agent learning (default behavior for RLlib's multi-agent
|
||||
setups), also only `compute_loss_for_module()` should be overridden, but it will
|
||||
be called for each individual RLModule inside the MultiRLModule.
|
||||
For the functional update to work, no `forward` call should be made
|
||||
within this method, especially not a non-functional one. Instead, use
|
||||
the model outputs provided by `fwd_out`.
|
||||
|
||||
Losses from `DifferentiableLearner` instances in `others_loss_per_module`
|
||||
can be leveraged for more advanced module-wise loss calculations.
|
||||
|
||||
Args:
|
||||
fwd_out: Output from a call to the `forward_train()` method of the
|
||||
underlying MultiRLModule (`self.module`) during training
|
||||
(`self.update()`).
|
||||
batch: The train batch that was used to compute `fwd_out`.
|
||||
others_loss_per_module: A list of losses per module id from the contained
|
||||
`DifferentiableLearner` instances in`self.others`.
|
||||
|
||||
Returns:
|
||||
A dictionary mapping module IDs to individual loss terms.
|
||||
"""
|
||||
loss_per_module = {}
|
||||
for module_id in fwd_out:
|
||||
module_batch = batch[module_id]
|
||||
module_fwd_out = fwd_out[module_id]
|
||||
|
||||
module = self.module[module_id].unwrapped()
|
||||
if isinstance(module, SelfSupervisedLossAPI):
|
||||
loss = module.compute_self_supervised_loss(
|
||||
learner=self,
|
||||
module_id=module_id,
|
||||
config=self.config.get_config_for_module(module_id),
|
||||
batch=module_batch,
|
||||
fwd_out=module_fwd_out,
|
||||
others_loss_per_module=others_loss_per_module,
|
||||
)
|
||||
else:
|
||||
loss = self.compute_loss_for_module(
|
||||
module_id=module_id,
|
||||
config=self.config.get_config_for_module(module_id),
|
||||
batch=module_batch,
|
||||
fwd_out=module_fwd_out,
|
||||
others_loss_per_module=others_loss_per_module,
|
||||
)
|
||||
loss_per_module[module_id] = loss
|
||||
|
||||
return loss_per_module
|
||||
|
||||
@OverrideToImplementCustomLogic
|
||||
@override(Learner)
|
||||
def compute_loss_for_module(
|
||||
self,
|
||||
*,
|
||||
module_id: ModuleID,
|
||||
config: "AlgorithmConfig",
|
||||
batch: Dict[str, Any],
|
||||
fwd_out: Dict[str, TensorType],
|
||||
others_loss_per_module: List[Dict[ModuleID, TensorType]] = None,
|
||||
) -> TensorType:
|
||||
"""Computes the loss for a single module.
|
||||
|
||||
This method calculates the loss for an individual agent. In multi-agent
|
||||
scenarios requiring more complex loss computations, consider overriding the
|
||||
`compute_losses` method instead.
|
||||
|
||||
Losses from `DifferentiableLearner` instances in `others_loss_per_module`
|
||||
can be leveraged for more advanced module-wise loss calculations.
|
||||
|
||||
Args:
|
||||
module_id: The id of the module.
|
||||
config: The AlgorithmConfig specific to the given `module_id`.
|
||||
batch: The train batch for this particular module.
|
||||
fwd_out: The output of the forward pass for this particular module.
|
||||
others_loss_per_module: A list of losses per module id from the contained
|
||||
`DifferentiableLearner` instances in`self.others`.
|
||||
|
||||
Returns:
|
||||
A single total loss tensor. If you have more than one optimizer on the
|
||||
provided `module_id` and would like to compute gradients separately using
|
||||
these different optimizers, simply add up the individual loss terms for
|
||||
each optimizer and return the sum. Also, for recording/logging any
|
||||
individual loss terms, you can use the `Learner.metrics.log_value(
|
||||
key=..., value=...)` or `Learner.metrics.log_dict()` APIs. See:
|
||||
:py:class:`~ray.rllib.utils.metrics.metrics_logger.MetricsLogger` for more
|
||||
information.
|
||||
"""
|
||||
# By default, `others_loss_per_module` is not used; instead, the superclass
|
||||
# method is called.
|
||||
return super().compute_loss_for_module(
|
||||
module_id=module_id, config=config, batch=batch, fwd_out=fwd_out
|
||||
)
|
||||
|
||||
def _make_functional_call(
|
||||
self, params: Dict[ModuleID, NamedParamDict], batch: MultiAgentBatch
|
||||
) -> Dict[ModuleID, NamedParamDict]:
|
||||
"""Make a functional forward call to all modules in the `MultiRLModule`."""
|
||||
return self._module.foreach_module(
|
||||
lambda mid, m: torch.func.functional_call(m, params[mid], batch[mid]),
|
||||
return_dict=True,
|
||||
)
|
||||
@@ -0,0 +1,159 @@
|
||||
import dataclasses
|
||||
from collections import defaultdict
|
||||
from typing import List, Optional
|
||||
|
||||
import tree # pip install dm_tree
|
||||
|
||||
import ray
|
||||
from ray.rllib.env.multi_agent_episode import MultiAgentEpisode
|
||||
from ray.rllib.policy.sample_batch import MultiAgentBatch
|
||||
from ray.rllib.utils.minibatch_utils import (
|
||||
ShardBatchIterator,
|
||||
ShardEpisodesIterator,
|
||||
ShardObjectRefIterator,
|
||||
)
|
||||
from ray.rllib.utils.typing import EpisodeType
|
||||
|
||||
|
||||
# TODO (sven): Switch to dataclass(slots=True) once on py >= 3.10.
|
||||
@dataclasses.dataclass
|
||||
class TrainingData:
|
||||
batch: Optional[MultiAgentBatch] = None
|
||||
batches: Optional[List[MultiAgentBatch]] = None
|
||||
batch_refs: Optional[List[ray.ObjectRef]] = None
|
||||
|
||||
episodes: Optional[List[EpisodeType]] = None
|
||||
episodes_refs: Optional[List[ray.ObjectRef]] = None
|
||||
|
||||
data_iterators: Optional[List[ray.data.iterator.DataIterator]] = None
|
||||
|
||||
def validate(self):
|
||||
# Exactly one training data type must be provided.
|
||||
if (
|
||||
sum(
|
||||
td is not None
|
||||
for td in [
|
||||
self.batch,
|
||||
self.batches,
|
||||
self.batch_refs,
|
||||
self.episodes,
|
||||
self.episodes_refs,
|
||||
self.data_iterators,
|
||||
]
|
||||
)
|
||||
!= 1
|
||||
):
|
||||
raise ValueError("Exactly one training data type must be provided!")
|
||||
|
||||
def shard(
|
||||
self,
|
||||
num_shards: int,
|
||||
len_lookback_buffer: Optional[int] = None,
|
||||
**kwargs,
|
||||
):
|
||||
# Single batch -> Split into n smaller batches.
|
||||
if self.batch is not None:
|
||||
return [
|
||||
(TrainingData(batch=b), {})
|
||||
for b in ShardBatchIterator(self.batch, num_shards=num_shards)
|
||||
]
|
||||
|
||||
# TODO (sven): Do we need a more sohpisticated shard mechanism for this case?
|
||||
elif self.batches is not None:
|
||||
assert num_shards == len(self.batches)
|
||||
return [(TrainingData(batch=b), {}) for b in self.batches]
|
||||
|
||||
# List of batch refs.
|
||||
elif self.batch_refs is not None:
|
||||
return [
|
||||
(TrainingData(batch_refs=b), {})
|
||||
for b in ShardObjectRefIterator(self.batch_refs, num_shards=num_shards)
|
||||
]
|
||||
|
||||
# List of episodes -> Split into n equally sized shards (based on the lengths
|
||||
# of the episodes).
|
||||
elif self.episodes is not None:
|
||||
num_total_minibatches = 0
|
||||
if "minibatch_size" in kwargs and num_shards > 1:
|
||||
num_total_minibatches = self._compute_num_total_minibatches(
|
||||
self.episodes,
|
||||
num_shards,
|
||||
kwargs["minibatch_size"],
|
||||
kwargs.get("num_epochs", 1),
|
||||
)
|
||||
return [
|
||||
(
|
||||
TrainingData(episodes=e),
|
||||
{"num_total_minibatches": num_total_minibatches},
|
||||
)
|
||||
for e in ShardEpisodesIterator(
|
||||
self.episodes,
|
||||
num_shards=num_shards,
|
||||
len_lookback_buffer=len_lookback_buffer,
|
||||
)
|
||||
]
|
||||
# List of episodes refs.
|
||||
elif self.episodes_refs is not None:
|
||||
return [
|
||||
(TrainingData(episodes_refs=e), {})
|
||||
for e in ShardObjectRefIterator(self.episodes_refs, num_shards)
|
||||
]
|
||||
# List of data iterators.
|
||||
else:
|
||||
assert self.data_iterators and len(self.data_iterators) == num_shards
|
||||
return [
|
||||
(TrainingData(data_iterators=[di]), {}) for di in self.data_iterators
|
||||
]
|
||||
|
||||
def solve_refs(self):
|
||||
# Batch references.
|
||||
if self.batch_refs is not None:
|
||||
# Solve the ray.ObjRefs.
|
||||
batches = tree.flatten(ray.get(self.batch_refs))
|
||||
# If only a single batch, set `self.batch`.
|
||||
if len(batches) == 1:
|
||||
self.batch = batches[0]
|
||||
# Otherwise, set `self.batches`.
|
||||
else:
|
||||
self.batches = batches
|
||||
# Empty `self.batch_refs`.
|
||||
self.batch_refs = None
|
||||
|
||||
# Episode references.
|
||||
elif self.episodes_refs is not None:
|
||||
# It's possible that individual refs are invalid due to the EnvRunner
|
||||
# that produced the ref has crashed or had its entire node go down.
|
||||
# In this case, try each ref individually and collect only valid results.
|
||||
try:
|
||||
episodes = tree.flatten(ray.get(self.episodes_refs))
|
||||
except ray.exceptions.OwnerDiedError:
|
||||
episode_refs = self.episodes_refs
|
||||
episodes = []
|
||||
for ref in episode_refs:
|
||||
try:
|
||||
episodes.extend(ray.get(ref))
|
||||
except ray.exceptions.OwnerDiedError as e:
|
||||
ray.logger.warning(
|
||||
f"episode-ref {ref} died and can't be collected with error: {e}. This can happen if an EnvRunner is lost (for example because of a node failure) and is not critical in such cases."
|
||||
)
|
||||
self.episodes = episodes
|
||||
self.episodes_refs = None
|
||||
|
||||
@staticmethod
|
||||
def _compute_num_total_minibatches(
|
||||
episodes,
|
||||
num_shards,
|
||||
minibatch_size,
|
||||
num_epochs,
|
||||
):
|
||||
# Count total number of timesteps per module ID.
|
||||
if isinstance(episodes[0], MultiAgentEpisode):
|
||||
per_mod_ts = defaultdict(int)
|
||||
for ma_episode in episodes:
|
||||
for sa_episode in ma_episode.agent_episodes.values():
|
||||
per_mod_ts[sa_episode.module_id] += len(sa_episode)
|
||||
max_ts = max(per_mod_ts.values())
|
||||
else:
|
||||
max_ts = sum(map(len, episodes))
|
||||
|
||||
return int((num_epochs * max_ts) / (num_shards * minibatch_size))
|
||||
@@ -0,0 +1,58 @@
|
||||
import copy
|
||||
|
||||
from ray.rllib.utils.framework import try_import_torch
|
||||
from ray.rllib.utils.typing import NetworkType
|
||||
from ray.util import PublicAPI
|
||||
|
||||
torch, _ = try_import_torch()
|
||||
|
||||
|
||||
def make_target_network(main_net: NetworkType) -> NetworkType:
|
||||
"""Creates a (deep) copy of `main_net` (including synched weights) and returns it.
|
||||
|
||||
Args:
|
||||
main_net: The main network to return a target network for
|
||||
|
||||
Returns:
|
||||
The copy of `main_net` that can be used as a target net. Note that the weights
|
||||
of the returned net are already synched (identical) with `main_net`.
|
||||
"""
|
||||
# Deepcopy the main net (this should already take care of synching all weights).
|
||||
target_net = copy.deepcopy(main_net)
|
||||
# Make the target net not trainable.
|
||||
if isinstance(main_net, torch.nn.Module):
|
||||
target_net.requires_grad_(False)
|
||||
else:
|
||||
raise ValueError(f"Unsupported framework for given `main_net` {main_net}!")
|
||||
|
||||
return target_net
|
||||
|
||||
|
||||
@PublicAPI(stability="beta")
|
||||
def update_target_network(
|
||||
*,
|
||||
main_net: NetworkType,
|
||||
target_net: NetworkType,
|
||||
tau: float,
|
||||
) -> None:
|
||||
"""Updates a target network (from a "main" network) using Polyak averaging.
|
||||
|
||||
Thereby:
|
||||
new_target_net_weight = (
|
||||
tau * main_net_weight + (1.0 - tau) * current_target_net_weight
|
||||
)
|
||||
|
||||
Args:
|
||||
main_net: The nn.Module to update from.
|
||||
target_net: The target network to update.
|
||||
tau: The tau value to use in the Polyak averaging formula. Use 1.0 for a
|
||||
complete sync of the weights (target and main net will be the exact same
|
||||
after updating).
|
||||
"""
|
||||
if isinstance(main_net, torch.nn.Module):
|
||||
from ray.rllib.utils.torch_utils import update_target_network as _update_target
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unsupported framework for given `main_net` {main_net}!")
|
||||
|
||||
_update_target(main_net=main_net, target_net=target_net, tau=tau)
|
||||
Reference in New Issue
Block a user