Simplify the class configurations using enums.

Author: ernestum

src/imitation_cli/algorithm_configurations/airl.py

@@ -24,7 +24,7 @@ class Config:
     reward_net: reward_network.Config = MISSING
     demo_batch_size: int = 64
     n_disc_updates_per_round: int = 2
-    disc_opt_cls: optimizer_class.Config = optimizer_class.Adam()
+    disc_opt_cls: optimizer_class.Config = optimizer_class.Adam
     gen_train_timesteps: Optional[int] = None
     gen_replay_buffer_capacity: Optional[int] = None
     init_tensorboard: bool = False

src/imitation_cli/utils/activation_function_class.py

@@ -1,59 +1,37 @@
 """Classes for configuring activation functions."""
 import dataclasses
+from enum import Enum
 
+import torch
 from hydra.core.config_store import ConfigStore
 
 
-@dataclasses.dataclass
-class Config:
-    """Base class for activation function configs."""
-
-    # Note: we don't define _target_ here so in the subclasses it can be defined last.
-    #  This is the same pattern we use as in schedule.py.
-    pass
-
-
-@dataclasses.dataclass
-class TanH(Config):
-    """Config for TanH activation function."""
-
-    _target_: str = "imitation_cli.utils.activation_function_class.TanH.make"
-
-    @staticmethod
-    def make() -> type:
-        import torch
+class ActivationFunctionClass(Enum):
+    """Enum of activation function classes."""
 
-        return torch.nn.Tanh
+    TanH = torch.nn.Tanh
+    ReLU = torch.nn.ReLU
+    LeakyReLU = torch.nn.LeakyReLU
 
 
 @dataclasses.dataclass
-class ReLU(Config):
-    """Config for ReLU activation function."""
+class Config:
+    """Base class for activation function configs."""
 
-    _target_: str = "imitation_cli.utils.activation_function_class.ReLU.make"
+    activation_function_class: ActivationFunctionClass
+    _target_: str = "imitation_cli.utils.activation_function_class.Config.make"
 
     @staticmethod
-    def make() -> type:
-        import torch
-
-        return torch.nn.ReLU
-
+    def make(activation_function_class: ActivationFunctionClass) -> type:
+        return activation_function_class.value
 
-@dataclasses.dataclass
-class LeakyReLU(Config):
-    """Config for LeakyReLU activation function."""
-
-    _target_: str = "imitation_cli.utils.activation_function_class.LeakyReLU.make"
-
-    @staticmethod
-    def make() -> type:
-        import torch
 
-        return torch.nn.LeakyReLU
+TanH = Config(ActivationFunctionClass.TanH)
+ReLU = Config(ActivationFunctionClass.ReLU)
+LeakyReLU = Config(ActivationFunctionClass.LeakyReLU)
 
 
 def register_configs(group: str):
     cs = ConfigStore.instance()
-    cs.store(group=group, name="tanh", node=TanH)
-    cs.store(group=group, name="relu", node=ReLU)
-    cs.store(group=group, name="leaky_relu", node=LeakyReLU)
+    for cls in ActivationFunctionClass:
+        cs.store(group=group, name=cls.name.lower(), node=Config(cls))

src/imitation_cli/utils/feature_extractor_class.py

@@ -1,46 +1,35 @@
 """Register Hydra configs for stable_baselines3 feature extractors."""
 import dataclasses
+from enum import Enum
 
+import stable_baselines3.common.torch_layers as torch_layers
 from hydra.core.config_store import ConfigStore
-from omegaconf import MISSING
 
 
-@dataclasses.dataclass
-class Config:
-    """Base config for stable_baselines3 feature extractors."""
+class FeatureExtractorClass(Enum):
+    """Enum of feature extractor classes."""
 
-    _target_: str = MISSING
+    FlattenExtractor = torch_layers.FlattenExtractor
+    NatureCNN = torch_layers.NatureCNN
 
 
 @dataclasses.dataclass
-class FlattenExtractorConfig(Config):
-    """Config for FlattenExtractor."""
+class Config:
+    """Base config for stable_baselines3 feature extractors."""
 
-    _target_: str = (
-        "imitation_cli.utils.feature_extractor_class.FlattenExtractorConfig.make"
-    )
+    feature_extractor_class: FeatureExtractorClass
+    _target_: str = "imitation_cli.utils.feature_extractor_class.Config.make"
 
     @staticmethod
-    def make() -> type:
-        import stable_baselines3
-
-        return stable_baselines3.common.torch_layers.FlattenExtractor
+    def make(feature_extractor_class: FeatureExtractorClass) -> type:
+        return feature_extractor_class.value
 
 
-@dataclasses.dataclass
-class NatureCNNConfig(Config):
-    """Config for NatureCNN."""
-
-    _target_: str = "imitation_cli.utils.feature_extractor_class.NatureCNNConfig.make"
-
-    @staticmethod
-    def make() -> type:
-        import stable_baselines3
-
-        return stable_baselines3.common.torch_layers.NatureCNN
+FlattenExtractor = Config(FeatureExtractorClass.FlattenExtractor)
+NatureCNN = Config(FeatureExtractorClass.NatureCNN)
 
 
 def register_configs(group: str):
     cs = ConfigStore.instance()
-    cs.store(group=group, name="flatten", node=FlattenExtractorConfig)
-    cs.store(group=group, name="nature_cnn", node=NatureCNNConfig)
+    for cls in FeatureExtractorClass:
+        cs.store(group=group, name=cls.name.lower(), node=Config(cls))

src/imitation_cli/utils/optimizer_class.py

@@ -1,44 +1,35 @@
 """Register optimizer classes with Hydra."""
 import dataclasses
+from enum import Enum
 
+import torch
 from hydra.core.config_store import ConfigStore
-from omegaconf import MISSING
 
 
-@dataclasses.dataclass
-class Config:
-    """Base config for optimizer classes."""
+class OptimizerClass(Enum):
+    """Enum of optimizer classes."""
 
-    _target_: str = MISSING
+    Adam = torch.optim.Adam
+    SGD = torch.optim.SGD
 
 
 @dataclasses.dataclass
-class Adam(Config):
-    """Config for Adam optimizer class."""
+class Config:
+    """Base config for optimizer classes."""
 
-    _target_: str = "imitation_cli.utils.optimizer_class.Adam.make"
+    optimizer_class: OptimizerClass
+    _target_: str = "imitation_cli.utils.optimizer_class.Config.make"
 
     @staticmethod
-    def make() -> type:
-        import torch
-
-        return torch.optim.Adam
+    def make(optimizer_class: OptimizerClass) -> type:
+        return optimizer_class.value
 
 
-@dataclasses.dataclass
-class SGD(Config):
-    """Config for SGD optimizer class."""
-
-    _target_: str = "imitation_cli.utils.optimizer_class.SGD.make"
-
-    @staticmethod
-    def make() -> type:
-        import torch
-
-        return torch.optim.SGD
+Adam = Config(OptimizerClass.Adam)
+SGD = Config(OptimizerClass.SGD)
 
 
 def register_configs(group: str):
     cs = ConfigStore.instance()
-    cs.store(group=group, name="adam", node=Adam)
-    cs.store(group=group, name="sgd", node=SGD)
+    for cls in OptimizerClass:
+        cs.store(group=group, name=cls.name.lower(), node=Config(cls))

src/imitation_cli/utils/policy.py

@@ -65,20 +65,20 @@ class ActorCriticPolicy(Config):
     _target_: str = "imitation_cli.utils.policy.ActorCriticPolicy.make"
     lr_schedule: schedule.Config = schedule.FixedSchedule(3e-4)
     net_arch: Optional[Dict[str, List[int]]] = None
-    activation_fn: act_fun_class_cfg.Config = act_fun_class_cfg.TanH()
+    activation_fn: act_fun_class_cfg.Config = act_fun_class_cfg.TanH
     ortho_init: bool = True
     use_sde: bool = False
     log_std_init: float = 0.0
     full_std: bool = True
     use_expln: bool = False
     squash_output: bool = False
     features_extractor_class: feature_extractor_class_cfg.Config = (
-        feature_extractor_class_cfg.FlattenExtractorConfig()
+        feature_extractor_class_cfg.FlattenExtractor
     )
     features_extractor_kwargs: Optional[Dict[str, Any]] = None
     share_features_extractor: bool = True
     normalize_images: bool = True
-    optimizer_class: optimizer_class_cfg.Config = optimizer_class_cfg.Adam()
+    optimizer_class: optimizer_class_cfg.Config = optimizer_class_cfg.Adam
     optimizer_kwargs: Optional[Dict[str, Any]] = None
 
     @staticmethod