Formatting, typing and documentation fixes. Also the implicit seed dependency was pulled out of the utils and made explicit in airl.py

Author: ernestum

src/imitation_cli/airl.py

@@ -14,8 +14,8 @@
 from imitation_cli.algorithm_configurations import airl as airl_cfg
 from imitation_cli.utils import environment as environment_cfg
 from imitation_cli.utils import (
-    policy,
     policy_evaluation,
+    randomness,
     reward_network,
     rl_algorithm,
     trajectories,
@@ -26,7 +26,7 @@
 class RunConfig:
     """Config for running AIRL."""
 
-    seed: int = 0
+    rng: randomness.Config = randomness.Config(seed=0)
     total_timesteps: int = int(1e6)
     checkpoint_interval: int = 0
 
@@ -39,11 +39,11 @@ class RunConfig:
 
 
 cs = ConfigStore.instance()
-environment_cfg.register_configs("environment")
-trajectories.register_configs("airl/demonstrations", "${environment}")
-rl_algorithm.register_configs("airl/gen_algo", "${environment}")
+environment_cfg.register_configs("environment", "${rng}")
+trajectories.register_configs("airl/demonstrations", "${environment}", "${rng}")
+rl_algorithm.register_configs("airl/gen_algo", "${environment}", "${rng.seed}")
 reward_network.register_configs("airl/reward_net", "${environment}")
-policy_evaluation.register_configs("evaluation", "${environment}")
+policy_evaluation.register_configs("evaluation", "${environment}", "${rng}")
 
 cs.store(
     name="airl_run_base",

src/imitation_cli/config/airl_optuna.yaml

@@ -22,4 +22,4 @@ hydra:
   sweeper:
     params:
       environment: cartpole,pendulum
-      airl/reward_net: basic,shaped,ensemble
+      airl/reward_net: basic,shaped,small_ensemble

src/imitation_cli/config/airl_run.yaml

@@ -15,4 +15,4 @@ checkpoint_interval: 1
 airl:
   demo_batch_size: 128
   demonstrations:
-    total_timesteps: 10
+    total_timesteps: 10

src/imitation_cli/config/airl_sweep_env_and_rewardnet.yaml

@@ -21,4 +21,4 @@ hydra:
   sweeper:
     params:
       environment: cartpole,pendulum
-      airl/reward_net: basic,shaped,ensemble
+      airl/reward_net: basic,shaped,small_ensemble

src/imitation_cli/utils/__init__.py

@@ -0,0 +1 @@
+"""Configurations to be used as ingredient to algorithm configurations."""

src/imitation_cli/utils/activation_function_class.py

@@ -1,17 +1,22 @@
+"""Classes for configuring activation functions."""
 import dataclasses
 
 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
@@ -23,6 +28,8 @@ def make() -> type:
 
 @dataclasses.dataclass
 class ReLU(Config):
+    """Config for ReLU activation function."""
+
     _target_: str = "imitation_cli.utils.activation_function_class.ReLU.make"
 
     @staticmethod
@@ -34,10 +41,12 @@ def make() -> type:
 
 @dataclasses.dataclass
 class LeakyReLU(Config):
+    """Config for LeakyReLU activation function."""
+
     _target_: str = "imitation_cli.utils.activation_function_class.LeakyReLU.make"
 
     @staticmethod
-    def make()  -> type:
+    def make() -> type:
         import torch
 
         return torch.nn.LeakyReLU

src/imitation_cli/utils/environment.py

@@ -1,7 +1,9 @@
+"""Configuration for Gym environments."""
 from __future__ import annotations
+
 import dataclasses
 import typing
-from typing import Optional
+from typing import Optional, Union, cast
 
 if typing.TYPE_CHECKING:
     from stable_baselines3.common.vec_env import VecEnv
@@ -15,18 +17,21 @@
 
 @dataclasses.dataclass
 class Config:
+    """Configuration for Gym environments."""
+
     _target_: str = "imitation_cli.utils.environment.Config.make"
     env_name: str = MISSING  # The environment to train on
     n_envs: int = 8  # number of environments in VecEnv
-    parallel: bool = False  # Use SubprocVecEnv rather than DummyVecEnv  TODO: when setting this to true this is really slow for some reason
+    # TODO: when setting this to true this is really slow for some reason
+    parallel: bool = False  # Use SubprocVecEnv rather than DummyVecEnv
     max_episode_steps: int = MISSING  # Set to positive int to limit episode horizons
     env_make_kwargs: dict = dataclasses.field(
-        default_factory=dict
+        default_factory=dict,
     )  # The kwargs passed to `spec.make`.
-    rng: randomness.Config = randomness.Config()
+    rng: randomness.Config = MISSING
 
     @staticmethod
-    def make(log_dir: Optional[str]=None, **kwargs) -> VecEnv:
+    def make(log_dir: Optional[str] = None, **kwargs) -> VecEnv:
         from imitation.util import util
 
         return util.make_vec_env(log_dir=log_dir, **kwargs)
@@ -38,13 +43,24 @@ def make_rollout_venv(environment_config: Config) -> VecEnv:
     return call(
         environment_config,
         log_dir=None,
-        post_wrappers=[lambda env, i: wrappers.RolloutInfoWrapper(env)]
+        post_wrappers=[lambda env, i: wrappers.RolloutInfoWrapper(env)],
     )
 
 
-def register_configs(group: str):
+def register_configs(
+    group: str,
+    default_rng: Union[randomness.Config, str] = MISSING,
+):
+    default_rng = cast(randomness.Config, default_rng)
     cs = ConfigStore.instance()
-    cs.store(group=group, name="gym_env", node=Config)
-    cs.store(group=group, name="cartpole", node=Config(env_name="CartPole-v0", max_episode_steps=500))
-    cs.store(group=group, name="pendulum", node=Config(env_name="Pendulum-v1", max_episode_steps=500))
-
+    cs.store(group=group, name="gym_env", node=Config(rng=default_rng))
+    cs.store(
+        group=group,
+        name="cartpole",
+        node=Config(env_name="CartPole-v0", max_episode_steps=500, rng=default_rng),
+    )
+    cs.store(
+        group=group,
+        name="pendulum",
+        node=Config(env_name="Pendulum-v1", max_episode_steps=500, rng=default_rng),
+    )

src/imitation_cli/utils/feature_extractor_class.py

@@ -1,3 +1,4 @@
+"""Register Hydra configs for stable_baselines3 feature extractors."""
 import dataclasses
 
 from hydra.core.config_store import ConfigStore
@@ -6,12 +7,18 @@
 
 @dataclasses.dataclass
 class Config:
+    """Base config for stable_baselines3 feature extractors."""
+
     _target_: str = MISSING
 
 
 @dataclasses.dataclass
 class FlattenExtractorConfig(Config):
-    _target_: str = "imitation_cli.utils.feature_extractor_class.FlattenExtractorConfig.make"
+    """Config for FlattenExtractor."""
+
+    _target_: str = (
+        "imitation_cli.utils.feature_extractor_class.FlattenExtractorConfig.make"
+    )
 
     @staticmethod
     def make() -> type:
@@ -22,6 +29,8 @@ def make() -> type:
 
 @dataclasses.dataclass
 class NatureCNNConfig(Config):
+    """Config for NatureCNN."""
+
     _target_: str = "imitation_cli.utils.feature_extractor_class.NatureCNNConfig.make"
 
     @staticmethod

src/imitation_cli/utils/optimizer_class.py

@@ -1,3 +1,4 @@
+"""Register optimizer classes with Hydra."""
 import dataclasses
 
 from hydra.core.config_store import ConfigStore
@@ -6,11 +7,15 @@
 
 @dataclasses.dataclass
 class Config:
+    """Base config for optimizer classes."""
+
     _target_: str = MISSING
 
 
 @dataclasses.dataclass
 class Adam(Config):
+    """Config for Adam optimizer class."""
+
     _target_: str = "imitation_cli.utils.optimizer_class.Adam.make"
 
     @staticmethod
@@ -22,6 +27,8 @@ def make() -> type:
 
 @dataclasses.dataclass
 class SGD(Config):
+    """Config for SGD optimizer class."""
+
     _target_: str = "imitation_cli.utils.optimizer_class.SGD.make"
 
     @staticmethod