[Feature] Enable loading Callbacks from hydra configs

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)

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 = {}

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

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()

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