Type A2C and PPO init

torchy_baselines/a2c/a2c.py

@@ -1,10 +1,15 @@
+from typing import List, Tuple, Type, Union, Callable, Optional, Dict, Any
+
 from gym import spaces
 import torch as th
 import torch.nn.functional as F
 
 from torchy_baselines.common.utils import explained_variance
-from torchy_baselines.ppo.ppo import PPO
 from torchy_baselines.common import logger
+from torchy_baselines.common.type_aliases import GymEnv
+from torchy_baselines.ppo.ppo import PPO
+from torchy_baselines.ppo.policies import PPOPolicy
+
 
 
 class A2C(PPO):
@@ -46,13 +51,27 @@ class A2C(PPO):
         Setting it to auto, the code will be run on the GPU if possible.
     :param _init_setup_model: (bool) Whether or not to build the network at the creation of the instance
     """
-    def __init__(self, policy, env, learning_rate=7e-4,
+    def __init__(self, policy: Union[str, Type[PPOPolicy]],
-                 n_steps=5, gamma=0.99, gae_lambda=1.0,
+                 env: Union[GymEnv, str],
-                 ent_coef=0.0, vf_coef=0.5, max_grad_norm=0.5,
+                 learning_rate: Union[float, Callable] = 7e-4,
-                 rms_prop_eps=1e-5, use_rms_prop=True, use_sde=False, sde_sample_freq=-1,
+                 n_steps: int = 5,
-                 normalize_advantage=False, tensorboard_log=None, create_eval_env=False,
+                 gamma: float = 0.99,
-                 policy_kwargs=None, verbose=0, seed=None, device='auto',
+                 gae_lambda: float = 1.0,
-                 _init_setup_model=True):
+                 ent_coef: float = 0.0,
+                 vf_coef: float = 0.5,
+                 max_grad_norm: float = 0.5,
+                 rms_prop_eps: float = 1e-5,
+                 use_rms_prop: bool = True,
+                 use_sde: bool = False,
+                 sde_sample_freq: int = -1,
+                 normalize_advantage: bool = False,
+                 tensorboard_log: Optional[str] = None,
+                 create_eval_env: bool = False,
+                 policy_kwargs: Optional[Dict[str, Any]] = None,
+                 verbose: int = 0,
+                 seed: Optional[int] = None,
+                 device: Union[th.device, str] = 'auto',
+                 _init_setup_model: bool = True):
 
         super(A2C, self).__init__(policy, env, learning_rate=learning_rate,
                                   n_steps=n_steps, batch_size=None, n_epochs=1,

torchy_baselines/ppo/ppo.py

@@ -1,6 +1,6 @@
 import os
 import time
-from typing import Optional, Tuple, List
+from typing import List, Tuple, Type, Union, Callable, Optional, Dict, Any
 
 import gym
 from gym import spaces
@@ -14,12 +14,13 @@ except ImportError:
     SummaryWriter = None
 import numpy as np
 
+from torchy_baselines.common import logger
 from torchy_baselines.common.base_class import BaseRLModel
+from torchy_baselines.common.type_aliases import GymEnv
 from torchy_baselines.common.buffers import RolloutBuffer
 from torchy_baselines.common.utils import explained_variance, get_schedule_fn
 from torchy_baselines.common.vec_env import VecEnv
 from torchy_baselines.common.callbacks import BaseCallback
-from torchy_baselines.common import logger
 from torchy_baselines.ppo.policies import PPOPolicy
 
 
@@ -73,14 +74,29 @@ class PPO(BaseRLModel):
     :param _init_setup_model: (bool) Whether or not to build the network at the creation of the instance
     """
 
-    def __init__(self, policy, env, learning_rate=3e-4,
+    def __init__(self, policy: Union[str, Type[PPOPolicy]],
-                 n_steps=2048, batch_size=64, n_epochs=10,
+                 env: Union[GymEnv, str],
-                 gamma=0.99, gae_lambda=0.95, clip_range=0.2, clip_range_vf=None,
+                 learning_rate: Union[float, Callable] = 3e-4,
-                 ent_coef=0.0, vf_coef=0.5, max_grad_norm=0.5,
+                 n_steps: int = 2048,
-                 use_sde=False, sde_sample_freq=-1,
+                 batch_size: Optional[int] = 64,
-                 target_kl=None, tensorboard_log=None, create_eval_env=False,
+                 n_epochs: int = 10,
-                 policy_kwargs=None, verbose=0, seed=None, device='auto',
+                 gamma: float = 0.99,
-                 _init_setup_model=True):
+                 gae_lambda: float = 0.95,
+                 clip_range: float = 0.2,
+                 clip_range_vf: Optional[float] = None,
+                 ent_coef: float = 0.0,
+                 vf_coef: float = 0.5,
+                 max_grad_norm: float = 0.5,
+                 use_sde: bool = False,
+                 sde_sample_freq: int = -1,
+                 target_kl: Optional[float] = None,
+                 tensorboard_log: Optional[str] = None,
+                 create_eval_env: bool = False,
+                 policy_kwargs: Optional[Dict[str, Any]] = None,
+                 verbose: int = 0,
+                 seed: Optional[int] = None,
+                 device: Union[th.device, str] = 'auto',
+                 _init_setup_model: bool = True):
 
         super(PPO, self).__init__(policy, env, PPOPolicy, policy_kwargs=policy_kwargs,
                                   verbose=verbose, device=device, use_sde=use_sde, sde_sample_freq=sde_sample_freq,