saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-07-25 19:48:33 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Reformated every file with PEP 8 errors

Browse source
This commit is contained in:
Noah Dormann 2019-11-21 13:01:03 +01:00
parent b20b70db48
commit a7655ca6e1
10 changed files with 62 additions and 65 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
docs/conf.py
View file

@ -16,12 +16,14 @@ import os
import sys
from unittest.mock import MagicMock
import torchy_baselines
# source code directory, relative to this file, for sphinx-autobuild
sys.path.insert(0, os.path.abspath('..'))
class Mock(MagicMock):
__subclasses__ = []
@classmethod
def __getattr__(cls, name):
return MagicMock()
@ -42,10 +44,6 @@ MOCK_MODULES = ['joblib', 'scipy', 'scipy.signal',
'gym.wrappers', 'gym.wrappers.monitoring', 'zmq']
sys.modules.update((mod_name, Mock()) for mod_name in MOCK_MODULES)
import torchy_baselines
# -- Project information -----------------------------------------------------
project = 'Torchy Baselines'
@ -57,7 +55,6 @@ version = 'master (' + torchy_baselines.__version__ + ' )'
# The full version, including alpha/beta/rc tags
release = torchy_baselines.__version__
# -- General configuration ---------------------------------------------------
# If your documentation needs a minimal Sphinx version, state it here.
@ -102,7 +99,6 @@ exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store']
# The name of the Pygments (syntax highlighting) style to use.
pygments_style = 'sphinx'
# -- Options for HTML output -------------------------------------------------
# The theme to use for HTML and HTML Help pages. See the documentation for
@ -121,6 +117,7 @@ html_logo = '_static/img/logo.png'
def setup(app):
app.add_stylesheet("css/baselines_theme.css")
# Theme options are theme-specific and customize the look and feel of a theme
# further. For a list of options available for each theme, see the
# documentation.
@ -148,7 +145,6 @@ html_static_path = ['_static']
# Output file base name for HTML help builder.
htmlhelp_basename = 'TorchyBaselinesdoc'
# -- Options for LaTeX output ------------------------------------------------
latex_elements = {
@ -177,7 +173,6 @@ latex_documents = [
'Torchy Baselines Contributors', 'manual'),
]
# -- Options for manual page output ------------------------------------------
# One entry per manual page. List of tuples
@ -187,7 +182,6 @@ man_pages = [
[author], 1)
]
# -- Options for Texinfo output ----------------------------------------------
# Grouping the document tree into Texinfo files. List of tuples

1
tests/test_custom_policy.py
View file

@ -5,6 +5,7 @@ import pytest
from torchy_baselines import PPO
@pytest.mark.parametrize('net_arch', [
[12, dict(vf=[16], pi=[8])],
[4],

4
tests/test_run.py
View file

@ -6,7 +6,6 @@ import numpy as np
from torchy_baselines import A2C, CEMRL, PPO, SAC, TD3
from torchy_baselines.common.noise import NormalActionNoise
action_noise = NormalActionNoise(np.zeros(1), 0.1 * np.ones(1))
@ -19,8 +18,6 @@ def test_td3():
os.remove("test_save.pth")
def test_a2c():
model = A2C('MlpPolicy', 'CartPole-v1', policy_kwargs=dict(net_arch=[16]), verbose=1, create_eval_env=True)
model.learn(total_timesteps=1000, eval_freq=500)
@ -47,6 +44,7 @@ def test_onpolicy(model_class, env_id):
# model.load("test_save")
# os.remove("test_save.pth")
def test_sac():
model = SAC('MlpPolicy', 'Pendulum-v0', policy_kwargs=dict(net_arch=[64, 64]),
learning_starts=100, verbose=1, create_eval_env=True, ent_coef='auto',

9
tests/test_vec_envs.py
View file

@ -59,8 +59,10 @@ class CustomGymEnv(gym.Env):
@pytest.mark.parametrize('vec_env_wrapper', VEC_ENV_WRAPPERS)
def test_vecenv_custom_calls(vec_env_class, vec_env_wrapper):
"""Test access to methods/attributes of vectorized environments"""
def make_env():
return CustomGymEnv(gym.spaces.Box(low=np.zeros(2), high=np.ones(2)))
vec_env = vec_env_class([make_env for _ in range(N_ENVS)])
if vec_env_wrapper is not None:
@ -92,7 +94,6 @@ def test_vecenv_custom_calls(vec_env_class, vec_env_wrapper):
assert (env_method_subset[1] == np.ones((1, 3))).all()
assert len(env_method_subset) == 2
# Test to change value for all the environments
setattr_result = vec_env.set_attr('current_step', 42, indices=None)
getattr_result = vec_env.get_attr('current_step')
@ -193,8 +194,10 @@ SPACES = collections.OrderedDict([
('continuous', gym.spaces.Box(low=np.zeros(2), high=np.ones(2))),
])
def check_vecenv_spaces(vec_env_class, space, obs_assert):
"""Helper method to check observation spaces in vectorized environments."""
def make_env():
return CustomGymEnv(space)
@ -228,6 +231,7 @@ def test_vecenv_single_space(vec_env_class, space):
class _UnorderedDictSpace(gym.spaces.Dict):
"""Like DictSpace, but returns an unordered dict when sampling."""
def sample(self):
return dict(super().sample())
@ -301,14 +305,17 @@ class CustomWrapperB(VecNormalize):
def name_test(self):
return self.__class__
class CustomWrapperBB(CustomWrapperB):
def __init__(self, venv):
CustomWrapperB.__init__(self, venv)
self.var_bb = 'bb'
def test_vecenv_wrapper_getattr():
def make_env():
return CustomGymEnv(gym.spaces.Box(low=np.zeros(2), high=np.ones(2)))
vec_env = DummyVecEnv([make_env for _ in range(N_ENVS)])
wrapped = CustomWrapperA(CustomWrapperBB(vec_env))
assert wrapped.var_a == 'a'

10
torchy_baselines/a2c/a2c.py
View file

@ -115,15 +115,15 @@ class A2C(PPO):
self.policy.optimizer.step()
# approx_kl_divs.append(th.mean(old_log_prob - log_prob).detach().cpu().numpy())
# print(explained_variance(self.rollout_buffer.returns.flatten().cpu().numpy(),
# self.rollout_buffer.values.flatten().cpu().numpy()))
# print(explained_variance(self.rollout_buffer.returns.flatten().cpu().numpy(),
# self.rollout_buffer.values.flatten().cpu().numpy()))
def learn(self, total_timesteps, callback=None, log_interval=100,
eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="A2C", reset_num_timesteps=True):
return super(A2C, self).learn(total_timesteps=total_timesteps, callback=callback, log_interval=log_interval,
eval_env=eval_env, eval_freq=eval_freq, n_eval_episodes=n_eval_episodes,
tb_log_name=tb_log_name, reset_num_timesteps=reset_num_timesteps)
eval_env=eval_env, eval_freq=eval_freq, n_eval_episodes=n_eval_episodes,
tb_log_name=tb_log_name, reset_num_timesteps=reset_num_timesteps)
def save(self, path):
if not path.endswith('.pth'):
@ -135,4 +135,4 @@ class A2C(PPO):
path += '.pth'
if env is not None:
pass
self.policy.load_state_dict(th.load(path))
self.policy.load_state_dict(th.load(path))

80
torchy_baselines/common/base_class.py
View file

@ -290,9 +290,10 @@ class BaseRLModel(object):
if 'policy_kwargs' in kwargs and kwargs['policy_kwargs'] != data['policy_kwargs']:
raise ValueError("The specified policy kwargs do not equal the stored policy kwargs."
"Stored kwargs: {}, specified kwargs: {}".format(data['policy_kwargs'],kwargs['policy_kwargs']))
"Stored kwargs: {}, specified kwargs: {}".format(data['policy_kwargs'],
kwargs['policy_kwargs']))
model = cls(policy=data["policy"],env=None, _init_setup_model=False)
model = cls(policy=data["policy"], env=None, _init_setup_model=False)
model.__dict__.update(data)
model.__dict__.update(kwargs)
model.set_env(env)
@ -300,40 +301,8 @@ class BaseRLModel(object):
return model
@staticmethod
def _save_to_file_zip(save_path, data=None, params=None):
"""Save model to a zip archive
:param save_path: (str or file-like) Where to store the model
:param data: (OrderedDict) Class parameters being stored
:param params: (OrderedDict) Model parameters being stored expexted to be state_dict
"""
# data/params can be None, so do not
# try to serialize them blindly
if data is not None:
serialized_data = data_to_json(data)
# Check postfix if save_path is a string
if isinstance(save_path, str):
_, ext = os.path.splitext(save_path)
if ext == "":
save_path += ".zip"
# Create a zip-archive and write our objects
# there. This works when save_path is either
# str or a file-like
with zipfile.ZipFile(save_path, "w") as file_:
# Do not try to save "None" elements
if data is not None:
file_.writestr("data",serialized_data)
if params is not None:
with file_.open('param.pth', mode="w") as param_file:
th.save(params,param_file)
@staticmethod
def _load_from_file(load_path, load_data = True):
def _load_from_file(load_path, load_data=True):
""" Load model data from a .zip archive
:param load_path: (str or file-like) Where to load the model from
@ -351,7 +320,7 @@ class BaseRLModel(object):
# Open the zip archive and load data
try:
with zipfile.ZipFile(load_path,"r") as file_:
with zipfile.ZipFile(load_path, "r") as file_:
namelist = file_.namelist()
# If data or parameters is not in the
# zip archive, assume they were stored
@ -378,12 +347,6 @@ class BaseRLModel(object):
return data, params
def set_random_seed(self, seed=0):
set_random_seed(seed, using_cuda=self.device == th.device('cuda'))
self.action_space.seed(seed)
@ -478,7 +441,7 @@ class BaseRLModel(object):
# Display training infos
if self.verbose >= 1 and log_interval is not None and (
episode_num + total_episodes) % log_interval == 0:
episode_num + total_episodes) % log_interval == 0:
fps = int(num_timesteps / (time.time() - self.start_time))
logger.logkv("episodes", episode_num + total_episodes)
# logger.logkv("mean 100 episode reward", mean_reward)
@ -495,3 +458,34 @@ class BaseRLModel(object):
mean_reward = np.mean(episode_rewards) if total_episodes > 0 else 0.0
return mean_reward, total_steps, total_episodes, obs
def _save_to_file_zip(save_path, data=None, params=None):
"""Save model to a zip archive
:param save_path: (str or file-like) Where to store the model
:param data: (OrderedDict) Class parameters being stored
:param params: (OrderedDict) Model parameters being stored expexted to be state_dict
"""
# data/params can be None, so do not
# try to serialize them blindly
if data is not None:
serialized_data = data_to_json(data)
# Check postfix if save_path is a string
if isinstance(save_path, str):
_, ext = os.path.splitext(save_path)
if ext == "":
save_path += ".zip"
# Create a zip-archive and write our objects
# there. This works when save_path is either
# str or a file-like
with zipfile.ZipFile(save_path, "w") as file_:
# Do not try to save "None" elements
if data is not None:
file_.writestr("data", serialized_data)
if params is not None:
with file_.open('param.pth', mode="w") as param_file:
th.save(params, param_file)

4
torchy_baselines/common/distributions.py
View file

@ -4,6 +4,7 @@ import torch.nn as nn
from torch.distributions import Normal, Categorical
from gym import spaces
class Distribution(object):
def __init__(self):
super(Distribution, self).__init__()
@ -97,7 +98,8 @@ class SquashedDiagGaussianDistribution(DiagGaussianDistribution):
self.gaussian_action = None
def proba_distribution(self, mean_actions, log_std, deterministic=False):
action, _ = super(SquashedDiagGaussianDistribution, self).proba_distribution(mean_actions, log_std, deterministic)
action, _ = super(SquashedDiagGaussianDistribution, self).proba_distribution(mean_actions, log_std,
deterministic)
return action, self
def mode(self):

2
torchy_baselines/common/vec_env/subproc_vec_env.py
View file

@ -42,6 +42,7 @@ def _worker(remote, parent_remote, env_fn_wrapper):
except EOFError:
break
def tile_images(img_nhwc):
"""
Tile N images into one big PxQ image
@ -68,7 +69,6 @@ def tile_images(img_nhwc):
return out_image
class SubprocVecEnv(VecEnv):
"""
Creates a multiprocess vectorized wrapper for multiple environments, distributing each environment to its own

2
torchy_baselines/ppo/ppo.py
View file

@ -321,7 +321,7 @@ class PPO(BaseRLModel):
params_to_save = self.get_policy_parameters()
self._save_to_file_zip(path, data=data, params=params_to_save)
_save_to_file_zip(path, data=data, params=params_to_save)
"""def load(self, path, env=None, **_kwargs):
if not path.endswith('.pth'):

3
torchy_baselines/td3/td3.py
View file

@ -148,7 +148,8 @@ class TD3(BaseRLModel):
for param, target_param in zip(self.critic.parameters(), self.critic_target.parameters()):
target_param.data.copy_(tau * param.data + (1 - tau) * target_param.data)
def train_actor(self, gradient_steps=1, batch_size=100, tau_actor=0.005, tau_critic=0.005, replay_data=None):
def train_actor(self, gradient_steps: object = 1, batch_size: object = 100, tau_actor: object = 0.005, tau_critic: object = 0.005,
replay_data: object = None) -> object:
# Update optimizer learning rate
self._update_learning_rate(self.actor.optimizer)