From 286d6f3df97ef8736f3520fa7ad1668887bfcc49 Mon Sep 17 00:00:00 2001 From: Matteo Bettini Date: Tue, 19 Sep 2023 11:57:05 +0100 Subject: [PATCH] [Examples] Update hydra examples Signed-off-by: Matteo Bettini --- benchmarl/hydra_run.py | 55 ++++++++++++++++-------------------- examples/simple_hydra_run.py | 23 +++++---------- examples/vmas_run.py | 50 ++++++++++++++++++++++++++++++++ 3 files changed, 81 insertions(+), 47 deletions(-) create mode 100644 examples/vmas_run.py diff --git a/benchmarl/hydra_run.py b/benchmarl/hydra_run.py index 465b761e..9398ed96 100644 --- a/benchmarl/hydra_run.py +++ b/benchmarl/hydra_run.py @@ -1,20 +1,18 @@ -import hydra -from hydra.core.hydra_config import HydraConfig +from experiment import ExperimentConfig from omegaconf import DictConfig, OmegaConf -from benchmarl.environments import task_config_registry +from benchmarl.algorithms.common import AlgorithmConfig +from benchmarl.environments import Task, task_config_registry from benchmarl.experiment import Experiment from benchmarl.models import model_config_registry from benchmarl.models.common import ModelConfig, parse_model_config, SequenceModelConfig -def load_experiment_from_hydra_config(cfg: DictConfig, task_name: str) -> Experiment: - algorithm_config = OmegaConf.to_object(cfg.algorithm) - experiment_config = OmegaConf.to_object(cfg.experiment) - task_config = task_config_registry[task_name].update_config( - OmegaConf.to_container(cfg.task, resolve=True) - ) - model_config = load_model_from_hydra_config(cfg.model) +def load_experiment_from_hydra(cfg: DictConfig, task_name: str) -> Experiment: + algorithm_config = load_algorithm_config_from_hydra(cfg.algorithm) + experiment_config = load_experiment_config_from_hydra(cfg.experiment) + task_config = load_task_config_from_hydra(cfg.task, task_name) + model_config = load_model_config_from_hydra(cfg.model) return Experiment( task=task_config, @@ -25,10 +23,24 @@ def load_experiment_from_hydra_config(cfg: DictConfig, task_name: str) -> Experi ) -def load_model_from_hydra_config(cfg: DictConfig) -> ModelConfig: +def load_task_config_from_hydra(cfg: DictConfig, task_name: str) -> Task: + return task_config_registry[task_name].update_config( + OmegaConf.to_container(cfg.task, resolve=True) + ) + + +def load_experiment_config_from_hydra(cfg: DictConfig) -> ExperimentConfig: + return OmegaConf.to_object(cfg) + + +def load_algorithm_config_from_hydra(cfg: DictConfig) -> AlgorithmConfig: + return OmegaConf.to_object(cfg) + + +def load_model_config_from_hydra(cfg: DictConfig) -> ModelConfig: if "layers" in cfg.keys(): model_configs = [ - load_model_from_hydra_config(cfg.layers[f"l{i}"]) + load_model_config_from_hydra(cfg.layers[f"l{i}"]) for i in range(1, len(cfg.layers) + 1) ] return SequenceModelConfig( @@ -39,22 +51,3 @@ def load_model_from_hydra_config(cfg: DictConfig) -> ModelConfig: return model_class( **parse_model_config(OmegaConf.to_container(cfg, resolve=True)) ) - - -@hydra.main(version_base=None, config_path="conf", config_name="config") -def hydra_experiment(cfg: DictConfig) -> None: - hydra_choices = HydraConfig.get().runtime.choices - task_name = hydra_choices.task - print(f"\nAlgorithm: {hydra_choices.algorithm}, Task: {task_name}") - print("\nLoaded config:\n") - print(OmegaConf.to_yaml(cfg)) - - experiment = load_experiment_from_hydra_config( - cfg, - task_name=task_name, - ) - experiment.run() - - -if __name__ == "__main__": - hydra_experiment() diff --git a/examples/simple_hydra_run.py b/examples/simple_hydra_run.py index aeaeaa1a..bc5a941e 100644 --- a/examples/simple_hydra_run.py +++ b/examples/simple_hydra_run.py @@ -1,8 +1,5 @@ import hydra - -from benchmarl.environments import task_config_registry -from benchmarl.experiment import Experiment -from benchmarl.hydra_run import load_model_from_hydra_config +from benchmarl.hydra_run import load_experiment_from_hydra from hydra.core.hydra_config import HydraConfig from omegaconf import DictConfig, OmegaConf @@ -11,20 +8,14 @@ def hydra_experiment(cfg: DictConfig) -> None: hydra_choices = HydraConfig.get().runtime.choices task_name = hydra_choices.task + print(f"\nAlgorithm: {hydra_choices.algorithm}, Task: {task_name}") + print("\nLoaded config:\n") + print(OmegaConf.to_yaml(cfg)) - algorithm_config = OmegaConf.to_object(cfg.algorithm) - experiment_config = OmegaConf.to_object(cfg.experiment) - task_config = task_config_registry[task_name].update_config(cfg.task) - model_config = load_model_from_hydra_config(cfg.model) - - experiment = Experiment( - task=task_config, - algorithm_config=algorithm_config, - model_config=model_config, - seed=cfg.seed, - config=experiment_config, + experiment = load_experiment_from_hydra( + cfg, + task_name=task_name, ) - experiment.run() diff --git a/examples/vmas_run.py b/examples/vmas_run.py new file mode 100644 index 00000000..21678c8a --- /dev/null +++ b/examples/vmas_run.py @@ -0,0 +1,50 @@ +import hydra +from benchmarl.experiment import Experiment + +from benchmarl.hydra_run import ( + load_algorithm_config_from_hydra, + load_experiment_config_from_hydra, + load_model_config_from_hydra, + load_task_config_from_hydra, +) +from hydra.core.hydra_config import HydraConfig +from omegaconf import DictConfig, OmegaConf + + +@hydra.main(version_base=None, config_path="../benchmarl/conf", config_name="config") +def hydra_experiment(cfg: DictConfig) -> None: + hydra_choices = HydraConfig.get().runtime.choices + task_name = hydra_choices.task + print(f"\nAlgorithm: {hydra_choices.algorithm}, Task: {task_name}") + + algorithm_config = load_algorithm_config_from_hydra(cfg.algorithm) + task_config = load_task_config_from_hydra(cfg.task, task_name) + model_config = load_model_config_from_hydra(cfg.model) + experiment_config = cfg.experiment + + # Hyperparameter changes for VMAS experiments + experiment_config.sampling_device = "cuda" + experiment_config.train_device = "cuda" + experiment_config.collected_frames_per_batch = 60_000 + experiment_config.n_envs_per_worker = 600 + experiment_config.on_policy_minibatch_size = 4096 + experiment_config.evaluation_episodes = 200 + experiment_config = load_experiment_config_from_hydra(cfg.experiment) + + print("\nLoaded config:\n") + print(OmegaConf.to_yaml(cfg)) + + experiment = Experiment( + task=task_config, + algorithm_config=algorithm_config, + model_config=model_config, + seed=cfg.seed, + config=experiment_config, + ) + experiment.run() + + +if __name__ == "__main__": + hydra_experiment() + # To reproduce the VMAS results launch this with + # python run.py algorithm=ippo "task=vmas/navigation,vmas/balance,vmas/sampling" "seed=0,1,2"