wehub-resource-sync
157 lines
5.7 KiB
Python
157 lines
5.7 KiB
Python
from typing import Any, Callable, Dict, List, Optional
|
|
|
|
from ray.rllib.callbacks.callbacks import RLlibCallback
|
|
from ray.rllib.utils import force_list
|
|
from ray.rllib.utils.annotations import OldAPIStack
|
|
|
|
|
|
def make_callback(
|
|
callback_name: str,
|
|
callbacks_objects: Optional[List[RLlibCallback]] = None,
|
|
callbacks_functions: Optional[List[Callable]] = None,
|
|
*,
|
|
args: List[Any] = None,
|
|
kwargs: Dict[str, Any] = None,
|
|
) -> None:
|
|
"""Calls an RLlibCallback method or a registered callback callable.
|
|
|
|
Args:
|
|
callback_name: The name of the callback method or key, for example:
|
|
"on_episode_start" or "on_train_result".
|
|
callbacks_objects: The RLlibCallback object or list of RLlibCallback objects
|
|
to call the `callback_name` method on (in the order they appear in the
|
|
list).
|
|
callbacks_functions: The callable or list of callables to call
|
|
(in the order they appear in the list).
|
|
args: Call args to pass to the method/callable calls.
|
|
kwargs: Call kwargs to pass to the method/callable calls.
|
|
"""
|
|
# Loop through all available RLlibCallback objects.
|
|
callbacks_objects = force_list(callbacks_objects)
|
|
for callback_obj in callbacks_objects:
|
|
getattr(callback_obj, callback_name)(*(args or ()), **(kwargs or {}))
|
|
|
|
# Loop through all available RLlibCallback objects.
|
|
callbacks_functions = force_list(callbacks_functions)
|
|
for callback_fn in callbacks_functions:
|
|
callback_fn(*(args or ()), **(kwargs or {}))
|
|
|
|
|
|
@OldAPIStack
|
|
def _make_multi_callbacks(callback_class_list):
|
|
class _MultiCallbacks(RLlibCallback):
|
|
IS_CALLBACK_CONTAINER = True
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self._callback_list = [
|
|
callback_class() for callback_class in callback_class_list
|
|
]
|
|
|
|
def on_algorithm_init(self, **kwargs) -> None:
|
|
for callback in self._callback_list:
|
|
callback.on_algorithm_init(**kwargs)
|
|
|
|
def on_workers_recreated(self, **kwargs) -> None:
|
|
for callback in self._callback_list:
|
|
callback.on_workers_recreated(**kwargs)
|
|
|
|
# Only on new API stack.
|
|
def on_env_runners_recreated(self, **kwargs) -> None:
|
|
pass
|
|
|
|
def on_offline_eval_runners_recreated(self, **kwargs) -> None:
|
|
pass
|
|
|
|
def on_checkpoint_loaded(self, **kwargs) -> None:
|
|
for callback in self._callback_list:
|
|
callback.on_checkpoint_loaded(**kwargs)
|
|
|
|
def on_create_policy(self, *, policy_id, policy) -> None:
|
|
for callback in self._callback_list:
|
|
callback.on_create_policy(policy_id=policy_id, policy=policy)
|
|
|
|
def on_environment_created(self, **kwargs) -> None:
|
|
for callback in self._callback_list:
|
|
callback.on_environment_created(**kwargs)
|
|
|
|
def on_sub_environment_created(self, **kwargs) -> None:
|
|
for callback in self._callback_list:
|
|
callback.on_sub_environment_created(**kwargs)
|
|
|
|
def on_episode_created(self, **kwargs) -> None:
|
|
for callback in self._callback_list:
|
|
callback.on_episode_created(**kwargs)
|
|
|
|
def on_episode_start(self, **kwargs) -> None:
|
|
for callback in self._callback_list:
|
|
callback.on_episode_start(**kwargs)
|
|
|
|
def on_episode_step(self, **kwargs) -> None:
|
|
for callback in self._callback_list:
|
|
callback.on_episode_step(**kwargs)
|
|
|
|
def on_episode_end(self, **kwargs) -> None:
|
|
for callback in self._callback_list:
|
|
callback.on_episode_end(**kwargs)
|
|
|
|
def on_evaluate_start(self, **kwargs) -> None:
|
|
for callback in self._callback_list:
|
|
callback.on_evaluate_start(**kwargs)
|
|
|
|
def on_evaluate_end(self, **kwargs) -> None:
|
|
for callback in self._callback_list:
|
|
callback.on_evaluate_end(**kwargs)
|
|
|
|
# TODO (simon, sven): Fix the test such that we can simply remove
|
|
# these.
|
|
def on_evaluate_offline_start(self, **kwargs):
|
|
for callback in self._callback_list:
|
|
callback.on_evaluate_offline_start(**kwargs)
|
|
|
|
def on_evaluate_offline_end(self, **kwargs):
|
|
for callback in self._callback_list:
|
|
callback.on_evaluate_offline_end(**kwargs)
|
|
|
|
def on_postprocess_trajectory(
|
|
self,
|
|
*,
|
|
worker,
|
|
episode,
|
|
agent_id,
|
|
policy_id,
|
|
policies,
|
|
postprocessed_batch,
|
|
original_batches,
|
|
**kwargs,
|
|
) -> None:
|
|
for callback in self._callback_list:
|
|
callback.on_postprocess_trajectory(
|
|
worker=worker,
|
|
episode=episode,
|
|
agent_id=agent_id,
|
|
policy_id=policy_id,
|
|
policies=policies,
|
|
postprocessed_batch=postprocessed_batch,
|
|
original_batches=original_batches,
|
|
**kwargs,
|
|
)
|
|
|
|
def on_sample_end(self, **kwargs) -> None:
|
|
for callback in self._callback_list:
|
|
callback.on_sample_end(**kwargs)
|
|
|
|
def on_learn_on_batch(
|
|
self, *, policy, train_batch, result: dict, **kwargs
|
|
) -> None:
|
|
for callback in self._callback_list:
|
|
callback.on_learn_on_batch(
|
|
policy=policy, train_batch=train_batch, result=result, **kwargs
|
|
)
|
|
|
|
def on_train_result(self, **kwargs) -> None:
|
|
for callback in self._callback_list:
|
|
callback.on_train_result(**kwargs)
|
|
|
|
return _MultiCallbacks
|