Fix typing errors and typos

Makefile

@@ -3,10 +3,10 @@ SHELL=/bin/bash
 pytest:
 	./scripts/run_tests.sh
 
-pytype:
+type:
 	pytype
 
-doc:
+docs:
 	cd docs && make html
 
 spelling:

docs/modules/ppo.rst

@@ -24,6 +24,7 @@ Notes
 - Original paper: https://arxiv.org/abs/1707.06347
 - Clear explanation of PPO on Arxiv Insights channel: https://www.youtube.com/watch?v=5P7I-xPq8u8
 - OpenAI blog post: https://blog.openai.com/openai-baselines-ppo/
+- Spinning Up guide: https://spinningup.openai.com/en/latest/algorithms/ppo.html
 
 
 Can I use?

docs/spelling_wordlist.txt

@@ -113,3 +113,4 @@ optimizers
 Deprecations
 forkserver
 cuda
+Polyak

torchy_baselines/cem_rl/cem_rl.py

@@ -33,7 +33,7 @@ class CEMRL(TD3):
     :param learning_starts: (int) how many steps of the model to collect transitions for before learning starts
     :param gamma: (float) the discount factor
     :param batch_size: (int) Minibatch size for each gradient update
-    :param tau: (float) the soft update coefficient ("polyak update" of the target networks, between 0 and 1)
+    :param tau: (float) the soft update coefficient ("Polyak update" of the target networks, between 0 and 1)
     :param action_noise: (ActionNoise) the action noise type. Cf common.noise for the different action noise type.
     :param target_policy_noise: (float) Standard deviation of Gaussian noise added to target policy
         (smoothing noise)
@@ -103,6 +103,7 @@ class CEMRL(TD3):
               eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="CEMRL", reset_num_timesteps=True):
 
         timesteps_since_eval, episode_num, evaluations, obs, eval_env = self._setup_learn(eval_env)
+        actor_steps = 0
 
         while self.num_timesteps < total_timesteps:
 

torchy_baselines/common/base_class.py

@@ -2,6 +2,8 @@ import time
 import os
 import io
 import zipfile
+import typing
+from typing import Union, Type, Optional
 from abc import ABC, abstractmethod
 from collections import deque
 
@@ -10,13 +12,16 @@ import torch as th
 import numpy as np
 
 from torchy_baselines.common import logger
-from torchy_baselines.common.policies import get_policy_from_name
+from torchy_baselines.common.policies import BasePolicy, get_policy_from_name
 from torchy_baselines.common.utils import set_random_seed, get_schedule_fn, update_learning_rate
 from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv, unwrap_vec_normalize, sync_envs_normalization
 from torchy_baselines.common.monitor import Monitor
 from torchy_baselines.common.evaluation import evaluate_policy
 from torchy_baselines.common.save_util import data_to_json, json_to_data
 
+if typing.TYPE_CHECKING:
+    from torchy_baselines.common.noise import ActionNoise
+
 
 class BaseRLModel(ABC):
     """
@@ -43,7 +48,7 @@ class BaseRLModel(ABC):
     :param sde_sample_freq: (int) Sample a new noise matrix every n steps when using SDE
         Default: -1 (only sample at the beginning of the rollout)
     """
-    def __init__(self, policy, env, policy_base, policy_kwargs=None,
+    def __init__(self, policy: Type[BasePolicy], env: Union[gym.Env, VecEnv, str], policy_base, policy_kwargs=None,
                  verbose=0, device='auto', support_multi_env=False,
                  create_eval_env=False, monitor_wrapper=True, seed=None,
                  use_sde=False, sde_sample_freq=-1):
@@ -59,7 +64,7 @@ class BaseRLModel(ABC):
         if verbose > 0:
             print(f"Using {self.device} device")
 
-        self.env = env
+        self.env = None  # type: Union[gym.Env, VecEnv]
         # get VecNormalize object if needed
         self._vec_normalize_env = unwrap_vec_normalize(env)
         self.verbose = verbose
@@ -71,7 +76,10 @@ class BaseRLModel(ABC):
         self.eval_env = None
         self.replay_buffer = None
         self.seed = seed
-        self.action_noise = None
+        self.action_noise = None  # type: ActionNoise
+        self.start_time = None
+        self.policy, self.actor = None, None
+        self.learning_rate = None
         # Used for SDE only
         self.rollout_data = None
         self.on_policy_exploration = False
@@ -407,7 +415,7 @@ class BaseRLModel(ABC):
 
         return data, params, opt_params
 
-    def set_random_seed(self, seed=None):
+    def set_random_seed(self, seed: Optional[int] = None):
         """
         Set the seed of the pseudo-random generators
         (python, numpy, pytorch, gym, action_space)
@@ -443,7 +451,7 @@ class BaseRLModel(ABC):
             eval_env.seed(self.seed)
 
         eval_env = self._get_eval_env(eval_env)
-        obs = self.env.reset()
+        obs = self.env.reset()  # type: Union[gym.Env, VecEnv]
         return timesteps_since_eval, episode_num, evaluations, obs, eval_env
 
     def _update_info_buffer(self, infos):

torchy_baselines/common/vec_env/base_vec_env.py

@@ -146,6 +146,17 @@ class VecEnv(ABC):
         """
         raise NotImplementedError()
 
+    def seed(self, seed, indices=None):
+        """
+        :param seed: (int or [int])
+        :param indices: ([int])
+        """
+        indices = self._get_indices(indices)
+        if not hasattr(seed, 'len'):
+            seed = [seed] * len(indices)
+        assert len(seed) == len(indices)
+        return [self.env_method('seed', seed[i], indices=i) for i in indices]
+
     @property
     def unwrapped(self):
         if isinstance(self, VecEnvWrapper):

torchy_baselines/common/vec_env/dummy_vec_env.py

@@ -3,7 +3,7 @@ from copy import deepcopy
 
 import numpy as np
 
-from torchy_baselines.common.vec_env import VecEnv
+from torchy_baselines.common.vec_env.base_vec_env import VecEnv
 from torchy_baselines.common.vec_env.util import copy_obs_dict, dict_to_obs, obs_space_info
 
 
@@ -54,17 +54,6 @@ class DummyVecEnv(VecEnv):
             self._save_obs(env_idx, obs)
         return self._obs_from_buf()
 
-    def seed(self, seed, indices=None):
-        """
-        :param seed: (int or [int])
-        :param indices: ([int])
-        """
-        indices = self._get_indices(indices)
-        if not hasattr(seed, 'len'):
-            seed = [seed] * len(indices)
-        assert len(seed) == len(indices)
-        return [self.envs[i].seed(seed[i]) for i in indices]
-
     def close(self):
         for env in self.envs:
             env.close()

torchy_baselines/common/vec_env/subproc_vec_env.py

@@ -4,7 +4,7 @@ from collections import OrderedDict
 import gym
 import numpy as np
 
-from torchy_baselines.common.vec_env import VecEnv, CloudpickleWrapper
+from torchy_baselines.common.vec_env.base_vec_env import VecEnv, CloudpickleWrapper
 
 
 def _worker(remote, parent_remote, env_fn_wrapper):
@@ -111,7 +111,7 @@ class SubprocVecEnv(VecEnv):
         for work_remote, remote, env_fn in zip(self.work_remotes, self.remotes, env_fns):
             args = (work_remote, remote, CloudpickleWrapper(env_fn))
             # daemon=True: if the main process crashes, we should not cause things to hang
-            process = ctx.Process(target=_worker, args=args, daemon=True)
+            process = ctx.Process(target=_worker, args=args, daemon=True)  # pytype:disable=attribute-error
             process.start()
             self.processes.append(process)
             work_remote.close()
@@ -187,17 +187,6 @@ class SubprocVecEnv(VecEnv):
         for remote in target_remotes:
             remote.recv()
 
-    def seed(self, seed, indices=None):
-        """
-        :param seed: (int or [int])
-        :param indices: ([int])
-        """
-        indices = self._get_indices(indices)
-        if not hasattr(seed, 'len'):
-            seed = [seed] * len(indices)
-        assert len(seed) == len(indices)
-        return [self.env_method('seed', seed[i], indices=i) for i in indices]
-
     def env_method(self, method_name, *method_args, indices=None, **method_kwargs):
         """Call instance methods of vectorized environments."""
         target_remotes = self._get_target_remotes(indices)

torchy_baselines/common/vec_env/vec_frame_stack.py

@@ -3,7 +3,7 @@ import warnings
 import numpy as np
 from gym import spaces
 
-from torchy_baselines.common.vec_env import VecEnvWrapper
+from torchy_baselines.common.vec_env.base_vec_env import VecEnvWrapper
 
 
 class VecFrameStack(VecEnvWrapper):

torchy_baselines/common/vec_env/vec_normalize.py

@@ -2,7 +2,7 @@ import pickle
 
 import numpy as np
 
-from torchy_baselines.common.vec_env import VecEnvWrapper
+from torchy_baselines.common.vec_env.base_vec_env import VecEnvWrapper
 from torchy_baselines.common.running_mean_std import RunningMeanStd
 
 

torchy_baselines/ppo/policies.py

@@ -76,18 +76,20 @@ class PPOPolicy(BasePolicy):
 
         self.sde_feature_extractor = None
         self.sde_net_arch = sde_net_arch
+        self.use_sde = use_sde
 
         # Action distribution
         self.action_dist = make_proba_distribution(action_space, use_sde=use_sde, dist_kwargs=dist_kwargs)
 
         self._build(learning_rate)
 
-    def reset_noise(self, n_envs=1):
+    def reset_noise(self, n_envs: int = 1):
         """
         Sample new weights for the exploration matrix.
 
         :param n_envs: (int)
         """
+        assert isinstance(self.action_dist, StateDependentNoiseDistribution), 'reset_noise() is only available when using SDE'
         self.action_dist.sample_weights(self.log_std, batch_size=n_envs)
 
     def _build(self, learning_rate):

torchy_baselines/ppo/ppo.py

@@ -25,7 +25,7 @@ class PPO(BaseRLModel):
     Proximal Policy Optimization algorithm (PPO) (clip version)
 
     Paper: https://arxiv.org/abs/1707.06347
-    Code: This implementation borrows code from OpenAI spinningup (https://github.com/openai/spinningup/)
+    Code: This implementation borrows code from OpenAI Spinning Up (https://github.com/openai/spinningup/)
     https://github.com/ikostrikov/pytorch-a2c-ppo-acktr-gail and
     and Stable Baselines (PPO2 from https://github.com/hill-a/stable-baselines)
 

torchy_baselines/sac/policies.py

@@ -94,6 +94,7 @@ class Actor(BaseNetwork):
 
         :return: (th.Tensor)
         """
+        assert isinstance(self.action_dist, StateDependentNoiseDistribution), 'get_std() is only available when using SDE'
         return self.action_dist.get_std(self.log_std)
 
     def reset_noise(self, batch_size=1):
@@ -102,6 +103,7 @@ class Actor(BaseNetwork):
 
         :param batch_size: (int)
         """
+        assert isinstance(self.action_dist, StateDependentNoiseDistribution), 'reset_noise() is only available when using SDE'
         self.action_dist.sample_weights(self.log_std, batch_size=batch_size)
 
     def _get_latent(self, obs):

torchy_baselines/sac/sac.py

@@ -39,7 +39,7 @@ class SAC(BaseRLModel):
     :param gradient_steps: (int) How many gradient update after each step
     :param n_episodes_rollout: (int) Update the model every `n_episodes_rollout` episodes.
         Note that this cannot be used at the same time as `train_freq`
-    :param target_entropy: (str or float) target entropy when learning ent_coef (ent_coef = 'auto')
+    :param target_entropy: (str or float) target entropy when learning `ent_coef` (`ent_coef = 'auto'`)
     :param action_noise: (ActionNoise) the action noise type (None by default), this can help
         for hard exploration problem. Cf common.noise for the different action noise type.
     :param gamma: (float) the discount factor
@@ -161,7 +161,7 @@ class SAC(BaseRLModel):
         """
         return self.unscale_action(self.select_action(observation))
 
-    def train(self, gradient_steps, batch_size=64):
+    def train(self, gradient_steps: int, batch_size: int = 64):
         # Update optimizers learning rate
         optimizers = [self.actor.optimizer, self.critic.optimizer]
         if self.ent_coef_optimizer is not None:
@@ -169,6 +169,9 @@ class SAC(BaseRLModel):
 
         self._update_learning_rate(optimizers)
 
+        ent_coef_loss, ent_coef = th.zeros(1), th.zeros(1)
+        actor_loss, critic_loss = th.zeros(1), th.zeros(1)
+
         for gradient_step in range(gradient_steps):
             # Sample replay buffer
             replay_data = self.replay_buffer.sample(batch_size, env=self._vec_normalize_env)

torchy_baselines/td3/td3.py

@@ -31,7 +31,7 @@ class TD3(BaseRLModel):
     :param gradient_steps: (int) How many gradient update after each step
     :param n_episodes_rollout: (int) Update the model every `n_episodes_rollout` episodes.
         Note that this cannot be used at the same time as `train_freq`
-    :param tau: (float) the soft update coefficient ("polyak update" of the target networks, between 0 and 1)
+    :param tau: (float) the soft update coefficient ("Polyak update" of the target networks, between 0 and 1)
     :param action_noise: (ActionNoise) the action noise type. Cf common.noise for the different action noise type.
     :param target_policy_noise: (float) Standard deviation of Gaussian noise added to target policy
         (smoothing noise)