diff --git a/Makefile b/Makefile index c3636cf..78b043f 100644 --- a/Makefile +++ b/Makefile @@ -3,10 +3,10 @@ SHELL=/bin/bash pytest: ./scripts/run_tests.sh -pytype: +type: pytype -doc: +docs: cd docs && make html spelling: diff --git a/docs/modules/ppo.rst b/docs/modules/ppo.rst index 5531d3a..9c9887d 100644 --- a/docs/modules/ppo.rst +++ b/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? diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index 696ae55..6ac1c76 100644 --- a/docs/spelling_wordlist.txt +++ b/docs/spelling_wordlist.txt @@ -113,3 +113,4 @@ optimizers Deprecations forkserver cuda +Polyak diff --git a/torchy_baselines/cem_rl/cem_rl.py b/torchy_baselines/cem_rl/cem_rl.py index 1b84018..8ae7bb4 100644 --- a/torchy_baselines/cem_rl/cem_rl.py +++ b/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: diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index 5c780c9..a745d4c 100644 --- a/torchy_baselines/common/base_class.py +++ b/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): diff --git a/torchy_baselines/common/vec_env/base_vec_env.py b/torchy_baselines/common/vec_env/base_vec_env.py index d239ddd..e29c35e 100644 --- a/torchy_baselines/common/vec_env/base_vec_env.py +++ b/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): diff --git a/torchy_baselines/common/vec_env/dummy_vec_env.py b/torchy_baselines/common/vec_env/dummy_vec_env.py index c55ca6a..2d0211e 100644 --- a/torchy_baselines/common/vec_env/dummy_vec_env.py +++ b/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() diff --git a/torchy_baselines/common/vec_env/subproc_vec_env.py b/torchy_baselines/common/vec_env/subproc_vec_env.py index 0d1ecde..9205683 100644 --- a/torchy_baselines/common/vec_env/subproc_vec_env.py +++ b/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) diff --git a/torchy_baselines/common/vec_env/vec_frame_stack.py b/torchy_baselines/common/vec_env/vec_frame_stack.py index 2610f42..562c525 100644 --- a/torchy_baselines/common/vec_env/vec_frame_stack.py +++ b/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): diff --git a/torchy_baselines/common/vec_env/vec_normalize.py b/torchy_baselines/common/vec_env/vec_normalize.py index 4f5cb11..ea94bc0 100644 --- a/torchy_baselines/common/vec_env/vec_normalize.py +++ b/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 diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py index 5dd746f..4421d50 100644 --- a/torchy_baselines/ppo/policies.py +++ b/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): diff --git a/torchy_baselines/ppo/ppo.py b/torchy_baselines/ppo/ppo.py index 8641552..77a47cf 100644 --- a/torchy_baselines/ppo/ppo.py +++ b/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) diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py index f40f298..3fe11c5 100644 --- a/torchy_baselines/sac/policies.py +++ b/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): diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py index 70ca621..1e682cf 100644 --- a/torchy_baselines/sac/sac.py +++ b/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) diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py index 8a19d8f..32cc1b2 100644 --- a/torchy_baselines/td3/td3.py +++ b/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)