From 73b6642f8c71bbde21acf7cd69a81fec4de0867c Mon Sep 17 00:00:00 2001 From: Yoann Poupart <66315201+Xmaster6y@users.noreply.github.com> Date: Wed, 17 Dec 2025 12:05:21 +0100 Subject: [PATCH 1/2] Initial setup --- benchmarl/__init__.py | 6 +++ benchmarl/callbacks/__init__.py | 13 ++++++ benchmarl/callbacks/common.py | 78 +++++++++++++++++++++++++++++++++ benchmarl/hydra_config.py | 16 ++++++- 4 files changed, 111 insertions(+), 2 deletions(-) create mode 100644 benchmarl/callbacks/__init__.py create mode 100644 benchmarl/callbacks/common.py diff --git a/benchmarl/__init__.py b/benchmarl/__init__.py index 301fff5b..eac6fde6 100644 --- a/benchmarl/__init__.py +++ b/benchmarl/__init__.py @@ -23,6 +23,7 @@ def _load_hydra_schemas(): from hydra.core.config_store import ConfigStore from benchmarl.algorithms import algorithm_config_registry + from benchmarl.callbacks import callback_config_registry from benchmarl.environments import _task_class_registry from benchmarl.experiment import ExperimentConfig @@ -33,6 +34,11 @@ def _load_hydra_schemas(): # Load algos schemas for algo_name, algo_schema in algorithm_config_registry.items(): cs.store(name=f"{algo_name}_config", group="algorithm", node=algo_schema) + # Load callback schemas + for callback_name, callback_schema in callback_config_registry.items(): + cs.store( + name=f"{callback_name}_config", group="callback", node=callback_schema + ) # Load task schemas for task_schema_name, task_schema in _task_class_registry.items(): cs.store(name=f"{task_schema_name}_config", group="task", node=task_schema) diff --git a/benchmarl/callbacks/__init__.py b/benchmarl/callbacks/__init__.py new file mode 100644 index 00000000..188d7b1a --- /dev/null +++ b/benchmarl/callbacks/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. +# + +from .common import CallbackConfig + +__all__ = [ + "CallbackConfig", +] + +callback_config_registry = {} diff --git a/benchmarl/callbacks/common.py b/benchmarl/callbacks/common.py new file mode 100644 index 00000000..35d3b337 --- /dev/null +++ b/benchmarl/callbacks/common.py @@ -0,0 +1,78 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. +# + +import pathlib + +from abc import abstractmethod +from dataclasses import dataclass +from typing import Any, Dict, Optional, Type + +from benchmarl.experiment import Callback +from benchmarl.utils import _read_yaml_config + + +@dataclass +class CallbackConfig: + """ + Dataclass representing a callback configuration. + This should be overridden by implemented callbacks. + Implementors should: + + 1. add configuration parameters for their callback + 2. implement all abstract methods + + """ + + def get_callback(self) -> Callback: + """ + Main function to turn the config into the associated callback + + Returns: the Callback + + """ + return self.associated_class()( + **self.__dict__, # Passes all the custom config parameters + ) + + @staticmethod + def _load_from_yaml(name: str) -> Dict[str, Any]: + yaml_path = ( + pathlib.Path(__file__).parent.parent + / "conf" + / "callbacks" + / f"{name.lower()}.yaml" + ) + return _read_yaml_config(str(yaml_path.resolve())) + + @classmethod + def get_from_yaml(cls, path: Optional[str] = None): + """ + Load the callback configuration from yaml + + Args: + path (str, optional): The full path of the yaml file to load from. + If None, it will default to + ``benchmarl/conf/callbacks/self.associated_class().__name__`` + + Returns: the loaded CallbackConfig + """ + + if path is None: + config = CallbackConfig._load_from_yaml( + name=cls.associated_class().__name__ + ) + + else: + config = _read_yaml_config(path) + return cls(**config) + + @staticmethod + @abstractmethod + def associated_class() -> Type[Callback]: + """ + The callback class associated to the config + """ + raise NotImplementedError diff --git a/benchmarl/hydra_config.py b/benchmarl/hydra_config.py index f6d84bdf..52484a32 100644 --- a/benchmarl/hydra_config.py +++ b/benchmarl/hydra_config.py @@ -6,11 +6,12 @@ import importlib from dataclasses import is_dataclass from pathlib import Path +from typing import List from benchmarl.algorithms.common import AlgorithmConfig from benchmarl.environments import task_config_registry, TaskClass from benchmarl.environments.common import _type_check_task_config -from benchmarl.experiment import Experiment, ExperimentConfig +from benchmarl.experiment import Callback, Experiment, ExperimentConfig from benchmarl.models import model_config_registry from benchmarl.models.common import ModelConfig, parse_model_config, SequenceModelConfig @@ -48,6 +49,7 @@ def load_experiment_from_hydra( task_config = load_task_config_from_hydra(cfg.task, task_name) model_config = load_model_config_from_hydra(cfg.model) critic_model_config = load_model_config_from_hydra(cfg.critic_model) + _callbacks = load_callbacks_from_hydra(getattr(cfg, "callbacks", None) or {}) return Experiment( task=task_config, @@ -56,7 +58,7 @@ def load_experiment_from_hydra( critic_model_config=critic_model_config, seed=cfg.seed, config=experiment_config, - callbacks=callbacks, + callbacks=_callbacks + [*callbacks], ) @@ -134,6 +136,16 @@ def load_model_config_from_hydra(cfg: DictConfig) -> ModelConfig: ) +def load_callbacks_from_hydra(cfg: DictConfig) -> List[Callback]: + """Returns a list of :class:`~benchmarl.callbacks.Callback` from hydra config. + + Args: + cfg (DictConfig): the callbacks config dictionary from hydra + + """ + return [OmegaConf.to_object(callback).get_callback() for callback in cfg.values()] + + def _find_hydra_folder(restore_file: str) -> str: """Given the restore file, look for the .hydra folder max three levels above it.""" current_folder = Path(restore_file).parent.resolve() From a9f4e02c1ec3436bba821509e8d9d6da9a64fc1d Mon Sep 17 00:00:00 2001 From: Yoann Poupart <66315201+Xmaster6y@users.noreply.github.com> Date: Wed, 17 Dec 2025 12:05:33 +0100 Subject: [PATCH 2/2] Tests --- test/test_callbacks.py | 56 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 56 insertions(+) create mode 100644 test/test_callbacks.py diff --git a/test/test_callbacks.py b/test/test_callbacks.py new file mode 100644 index 00000000..6fed2e26 --- /dev/null +++ b/test/test_callbacks.py @@ -0,0 +1,56 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. +# + +import pytest + +from benchmarl.callbacks import callback_config_registry +from benchmarl.hydra_config import load_callbacks_from_hydra + +from hydra import compose, initialize + + +def test_no_callbacks(): + with initialize(version_base=None, config_path="../benchmarl/conf"): + cfg = compose( + config_name="config", + overrides=[ + "algorithm=mappo", + "task=vmas/balance", + ], + ) + callbacks = load_callbacks_from_hydra(getattr(cfg, "callbacks", None) or {}) + assert len(callbacks) == 0 + + +@pytest.mark.parametrize("callback_name", callback_config_registry.keys()) +def test_loading_callbacks(callback_name): + with initialize(version_base=None, config_path="../benchmarl/conf"): + cfg = compose( + config_name="config", + overrides=[ + "algorithm=mappo", + "task=vmas/balance", + f"callback@callbacks.c1={callback_name}", + ], + ) + callback = load_callbacks_from_hydra(cfg.callbacks)[0] + assert isinstance( + callback, callback_config_registry[callback_name].associated_class() + ) + + +def test_disabling_callbacks(): + with initialize(version_base=None, config_path="../benchmarl/conf"): + cfg = compose( + config_name="config", + overrides=[ + "algorithm=mappo", + "task=vmas/balance", + "+callbacks=null", + ], + ) + callbacks = load_callbacks_from_hydra(getattr(cfg, "callbacks", None) or {}) + assert len(callbacks) == 0