chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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
|
||||
Reference in New Issue
Block a user