saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-09-15 22:10:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Reformat

Browse source
This commit is contained in:
Antonin Raffin 2020-03-12 11:12:10 +01:00
parent 037986a91d
commit 18f38f8cf5
24 changed files with 67 additions and 61 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
tests/test_callbacks.py
View file

@ -6,7 +6,7 @@ import gym
from torchy_baselines import A2C, CEMRL, PPO, SAC, TD3
from torchy_baselines.common.callbacks import (CallbackList, CheckpointCallback, EvalCallback,
EveryNTimesteps, StopTrainingOnRewardThreshold)
EveryNTimesteps, StopTrainingOnRewardThreshold)
@pytest.mark.parametrize("model_class", [A2C, CEMRL, PPO, SAC, TD3])
@ -44,6 +44,6 @@ def test_callbacks(model_class):
# Transform callback into a callback list automatically
model.learn(500, callback=[checkpoint_callback, eval_callback])
# Automatic wrapping, old way of doing callbacks
model.learn(500, callback=lambda _locals, _globals : True)
model.learn(500, callback=lambda _locals, _globals: True)
if os.path.exists(log_folder):
shutil.rmtree(log_folder)

1
tests/test_distributions.py
View file

@ -22,6 +22,7 @@ def test_bijector():
# Check the inverse method
assert th.isclose(TanhBijector.inverse(squashed_actions), actions).all()
@pytest.mark.parametrize("model_class", [A2C, PPO])
def test_squashed_gaussian(model_class):
"""

3
tests/test_logger.py
View file

@ -5,7 +5,8 @@ import pytest
import numpy as np
from torchy_baselines.common.logger import (make_output_format, read_csv, read_json, DEBUG, ScopedConfigure,
info, debug, set_level, configure, logkv, logkvs, dumpkvs, logkv_mean, warn, error, reset)
info, debug, set_level, configure, logkv, logkvs, dumpkvs, logkv_mean, warn,
error, reset)
KEY_VALUES = {
"test": 1,

1
tests/test_predict.py
View file

@ -12,6 +12,7 @@ MODEL_LIST = [
SAC,
]
@pytest.mark.parametrize("model_class", MODEL_LIST)
def test_auto_wrap(model_class):
# test auto wrapping of env into a VecEnv

4
tests/test_run.py
View file

@ -1,7 +1,5 @@
import os
import pytest
import numpy as np
import pytest
from torchy_baselines import A2C, CEMRL, PPO, SAC, TD3
from torchy_baselines.common.noise import NormalActionNoise, OrnsteinUhlenbeckActionNoise

1
tests/test_save_load.py
View file

@ -16,6 +16,7 @@ MODEL_LIST = [
SAC,
]
@pytest.mark.parametrize("model_class", MODEL_LIST)
def test_save_load(model_class):
"""

7
tests/test_vec_normalize.py
View file

@ -8,9 +8,11 @@ from torchy_baselines import CEMRL, SAC, TD3
ENV_ID = 'Pendulum-v0'
def make_env():
return gym.make(ENV_ID)
def check_rms_equal(rmsa, rmsb):
assert np.all(rmsa.mean == rmsb.mean)
assert np.all(rmsa.var == rmsb.var)
@ -34,6 +36,7 @@ def check_vec_norm_equal(norma, normb):
assert norma.epsilon == normb.epsilon
assert norma.training == normb.training
def _make_warmstart_cartpole():
"""Warm-start VecNormalize by stepping through CartPole"""
venv = DummyVecEnv([lambda: gym.make("CartPole-v1")])
@ -50,8 +53,8 @@ def _make_warmstart_cartpole():
def test_runningmeanstd():
"""Test RunningMeanStd object"""
for (x_1, x_2, x_3) in [
(np.random.randn(3), np.random.randn(4), np.random.randn(5)),
(np.random.randn(3, 2), np.random.randn(4, 2), np.random.randn(5, 2))]:
(np.random.randn(3), np.random.randn(4), np.random.randn(5)),
(np.random.randn(3, 2), np.random.randn(4, 2), np.random.randn(5, 2))]:
rms = RunningMeanStd(epsilon=0.0, shape=x_1.shape[1:])
x_cat = np.concatenate([x_1, x_2, x_3], axis=0)

11
torchy_baselines/a2c/a2c.py
View file

@ -1,15 +1,14 @@
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 gym import spaces
from typing import Type, Union, Callable, Optional, Dict, Any
from torchy_baselines.common.utils import explained_variance
from torchy_baselines.common import logger
from torchy_baselines.common.type_aliases import GymEnv
from torchy_baselines.common.callbacks import BaseCallback
from torchy_baselines.ppo.ppo import PPO
from torchy_baselines.common.type_aliases import GymEnv
from torchy_baselines.common.utils import explained_variance
from torchy_baselines.ppo.policies import PPOPolicy
from torchy_baselines.ppo.ppo import PPO
class A2C(PPO):

3
torchy_baselines/cem_rl/cem.py
View file

@ -1,6 +1,5 @@
from typing import Type, Tuple, Optional, List
import numpy as np
from typing import Tuple, Optional, List
# TODO: add more from https://github.com/hardmaru/estool/blob/master/es.py

4
torchy_baselines/common/base_class.py
View file

@ -16,7 +16,7 @@ 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, VecNormalize
from torchy_baselines.common.save_util import data_to_json, json_to_data, recursive_getattr, recursive_setattr
from torchy_baselines.common.type_aliases import GymEnv, TensorDict, OptimizerStateDict, RolloutReturn
from torchy_baselines.common.type_aliases import GymEnv, TensorDict, RolloutReturn
from torchy_baselines.common.callbacks import BaseCallback, CallbackList, ConvertCallback, EvalCallback
from torchy_baselines.common.monitor import Monitor
from torchy_baselines.common.noise import ActionNoise
@ -494,7 +494,7 @@ class BaseRLModel(ABC):
if "data" in namelist and load_data:
# Load class parameters and convert to string
json_data = archive.read("data").decode()
data = json_to_data(json_data, device)
data = json_to_data(json_data)
if "tensors.pth" in namelist and load_data:
# Load extra tensors

10
torchy_baselines/common/buffers.py
View file

@ -18,6 +18,7 @@ class BaseBuffer(object):
to which the values will be converted
:param n_envs: (int) Number of parallel environments
"""
def __init__(self,
buffer_size: int,
obs_dim: int,
@ -118,13 +119,13 @@ class BaseBuffer(object):
@staticmethod
def _normalize_obs(obs: np.ndarray,
env: Optional[VecNormalize] = None) -> np.ndarray:
env: Optional[VecNormalize] = None) -> np.ndarray:
if env is not None:
return env.normalize_obs(obs).astype(np.float32)
return obs
def _normalize_reward(self,
reward: np.ndarray,
@staticmethod
def _normalize_reward(reward: np.ndarray,
env: Optional[VecNormalize] = None) -> np.ndarray:
if env is not None:
return env.normalize_reward(reward).astype(np.float32)
@ -141,13 +142,13 @@ class ReplayBuffer(BaseBuffer):
:param device: (th.device)
:param n_envs: (int) Number of parallel environments
"""
def __init__(self,
buffer_size: int,
obs_dim: int,
action_dim: int,
device: Union[th.device, str] = 'cpu',
n_envs: int = 1):
super(ReplayBuffer, self).__init__(buffer_size, obs_dim, action_dim, device, n_envs=n_envs)
assert n_envs == 1, "Replay buffer only support single environment for now"
@ -201,6 +202,7 @@ class RolloutBuffer(BaseBuffer):
:param gamma: (float) Discount factor
:param n_envs: (int) Number of parallel environments
"""
def __init__(self,
buffer_size: int,
obs_dim: int,

4
torchy_baselines/common/distributions.py
View file

@ -317,7 +317,7 @@ class StateDependentNoiseDistribution(Distribution):
self.exploration_matrices = self.weights_dist.rsample((batch_size,))
def proba_distribution_net(self, latent_dim: int, log_std_init: float = -2.0,
latent_sde_dim: Optional[th.Tensor] = None) -> Tuple[nn.Module, nn.Parameter]:
latent_sde_dim: Optional[int] = None) -> Tuple[nn.Module, nn.Parameter]:
"""
Create the layers and parameter that represent the distribution:
one output will be the deterministic action, the other parameter will be the
@ -325,7 +325,7 @@ class StateDependentNoiseDistribution(Distribution):
:param latent_dim: (int) Dimension of the last layer of the policy (before the action layer)
:param log_std_init: (float) Initial value for the log standard deviation
:param latent_sde_dim: (int) Dimension of the last layer of the feature extractor
:param latent_sde_dim: (Optional[int]) Dimension of the last layer of the feature extractor
for SDE. By default, it is shared with the policy network.
:return: (nn.Linear, nn.Parameter)
"""

5
torchy_baselines/common/monitor.py
View file

@ -27,8 +27,9 @@ class Monitor(gym.Wrapper):
:param env: (gym.Env) The environment
:param filename: (Optional[str]) the location to save a log file, can be None for no log
:param allow_early_resets: (bool) allows the reset of the environment before it is done
:param reset_keywords: (Tuple[str, ...]) extra keywords for the reset call, if extra parameters are needed at reset
:param info_keywords: (Tuple[str, ...]) extra information to log, from the information return of environment.step
:param reset_keywords: (Tuple[str, ...]) extra keywords for the reset call,
if extra parameters are needed at reset
:param info_keywords: (Tuple[str, ...]) extra information to log, from the information return of env.step()
"""
super(Monitor, self).__init__(env=env)
self.t_start = time.time()

3
torchy_baselines/common/noise.py
View file

@ -9,6 +9,7 @@ class ActionNoise(ABC):
"""
The action noise base class
"""
def __init__(self):
super(ActionNoise, self).__init__()
@ -22,6 +23,7 @@ class ActionNoise(ABC):
def __call__(self):
pass
class NormalActionNoise(ActionNoise):
"""
A Gaussian action noise
@ -29,6 +31,7 @@ class NormalActionNoise(ActionNoise):
:param mean: (float) the mean value of the noise
:param sigma: (float) the scale of the noise (std here)
"""
def __init__(self, mean, sigma):
self._mu = mean
self._sigma = sigma

2
torchy_baselines/common/save_util.py
View file

@ -122,14 +122,12 @@ def data_to_json(data: Dict[str, Any]) -> str:
def json_to_data(json_string: str,
device: Union[th.device, str] = 'cpu',
custom_objects: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
"""
Turn JSON serialization of class-parameters back into dictionary.
:param json_string: (str) JSON serialization of the class-parameters
that should be loaded.
:param device: torch.device device to which the data should be mapped if errors occur
:param custom_objects: (dict) Dictionary of objects to replace
upon loading. If a variable is present in this dictionary as a
key, it will not be deserialized and the corresponding item

5
torchy_baselines/common/type_aliases.py
View file

@ -1,8 +1,7 @@
"""
Common aliases for type hing
Common aliases for type hint
"""
from typing import Union, Type, Optional, Dict, Any, List, NamedTuple
from collections import namedtuple
from typing import Union, Dict, Any, NamedTuple, Optional
import numpy as np
import torch as th

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

@ -244,7 +244,7 @@ class VecEnvWrapper(VecEnv):
if blocked_class is not None:
own_class = f"{type(self).__module__}.{type(self).__name__}"
error_str = (f"Error: Recursive attribute lookup for {name} from {own_class} is "
"ambiguous and hides attribute from {blocked_class}")
"ambiguous and hides attribute from {blocked_class}")
raise AttributeError(error_str)
return self.getattr_recursive(name)

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

@ -61,11 +61,11 @@ def tile_images(img_nhwc):
new_width = int(np.ceil(float(n_images) / new_height))
img_nhwc = np.array(list(img_nhwc) + [img_nhwc[0] * 0 for _ in range(n_images, new_height * new_width)])
# img_HWhwc
out_image = img_nhwc.reshape(new_height, new_width, height, width, n_channels)
out_image = img_nhwc.reshape((new_height, new_width, height, width, n_channels))
# img_HhWwc
out_image = out_image.transpose(0, 2, 1, 3, 4)
# img_Hh_Ww_c
out_image = out_image.reshape(new_height * height, new_width * width, n_channels)
out_image = out_image.reshape((new_height * height, new_width * width, n_channels))
return out_image

6
torchy_baselines/common/vec_env/vec_normalize.py
View file

@ -86,7 +86,7 @@ class VecNormalize(VecEnvWrapper):
"""
obs, rews, news, infos = self.venv.step_wait()
self.old_obs = obs
self.old_rews = rews
self.old_reward = rews
if self.training:
self.obs_rms.update(obs)
@ -122,7 +122,7 @@ class VecNormalize(VecEnvWrapper):
"""
if self.norm_reward:
reward = np.clip(reward / np.sqrt(self.ret_rms.var + self.epsilon),
-self.clip_reward, self.clip_reward)
-self.clip_reward, self.clip_reward)
return reward
def unnormalize_obs(self, obs):
@ -146,7 +146,7 @@ class VecNormalize(VecEnvWrapper):
"""
Returns an unnormalized version of the rewards from the most recent step.
"""
return self.old_rews.copy()
return self.old_reward.copy()
def reset(self):
"""

10
torchy_baselines/ppo/policies.py
View file

@ -9,8 +9,8 @@ import numpy as np
from torchy_baselines.common.policies import (BasePolicy, register_policy, MlpExtractor,
create_sde_feature_extractor)
from torchy_baselines.common.distributions import (make_proba_distribution, Distribution,
DiagGaussianDistribution, CategoricalDistribution, StateDependentNoiseDistribution)
DiagGaussianDistribution, CategoricalDistribution,
StateDependentNoiseDistribution)
class PPOPolicy(BasePolicy):
@ -183,13 +183,15 @@ class PPOPolicy(BasePolicy):
action, _ = self._get_action_dist_from_latent(latent_pi, latent_sde, deterministic=deterministic)
return action
def evaluate_actions(self, obs: th.Tensor, actions: th.Tensor, deterministic: bool = False) -> Tuple[th.Tensor, th.Tensor, th.Tensor]:
def evaluate_actions(self, obs: th.Tensor,
actions: th.Tensor,
deterministic: bool = False) -> Tuple[th.Tensor, th.Tensor, th.Tensor]:
"""
Evaluate actions according to the current policy,
given the observations.
:param obs: (th.Tensor)
:param action: (th.Tensor)
:param actions: (th.Tensor)
:param deterministic: (bool)
:return: (th.Tensor, th.Tensor, th.Tensor) estimated value, log likelihood of taking those actions
and entropy of the action distribution.

18
torchy_baselines/ppo/ppo.py
View file

@ -146,11 +146,11 @@ class PPO(BaseRLModel):
self.clip_range_vf = get_schedule_fn(self.clip_range_vf)
def collect_rollouts(self,
env: VecEnv,
callback: BaseCallback,
rollout_buffer: RolloutBuffer,
n_rollout_steps: int = 256,
obs: Optional[np.ndarray] = None) -> Tuple[Optional[np.ndarray], bool]:
env: VecEnv,
callback: BaseCallback,
rollout_buffer: RolloutBuffer,
n_rollout_steps: int = 256,
obs: Optional[np.ndarray] = None) -> Tuple[Optional[np.ndarray], bool]:
n_steps = 0
continue_training = True
@ -167,7 +167,6 @@ class PPO(BaseRLModel):
continue_training = False
return None, continue_training
if self.use_sde and self.sde_sample_freq > 0 and n_steps % self.sde_sample_freq == 0:
# Sample a new noise matrix
self.policy.reset_noise(env.num_envs)
@ -227,7 +226,8 @@ class PPO(BaseRLModel):
values, log_prob, entropy = self.policy.evaluate_actions(rollout_data.observations, actions)
values = values.flatten()
# Normalize advantage
advantages = (rollout_data.advantages - rollout_data.advantages.mean()) / (rollout_data.advantages.std() + 1e-8)
advantages = (rollout_data.advantages - rollout_data.advantages.mean()) / (
rollout_data.advantages.std() + 1e-8)
# ratio between old and new policy, should be one at the first iteration
ratio = th.exp(log_prob - rollout_data.old_log_prob)
@ -242,7 +242,8 @@ class PPO(BaseRLModel):
else:
# Clip the different between old and new value
# NOTE: this depends on the reward scaling
values_pred = rollout_data.old_values + th.clamp(values - rollout_data.old_values, -clip_range_vf, clip_range_vf)
values_pred = rollout_data.old_values + th.clamp(values - rollout_data.old_values, -clip_range_vf,
clip_range_vf)
# Value loss using the TD(gae_lambda) target
value_loss = F.mse_loss(rollout_data.returns, values_pred)
@ -275,7 +276,6 @@ class PPO(BaseRLModel):
if self.clip_range_vf is not None:
logger.logkv("clip_range_vf", clip_range_vf)
logger.logkv("explained_variance", explained_var)
# TODO: gather stats for the entropy and other losses?
logger.logkv("entropy_loss", entropy_loss.item())

1
torchy_baselines/sac/policies.py
View file

@ -244,6 +244,7 @@ class SACPolicy(BasePolicy):
def predict(self, observation: th.Tensor, deterministic: bool = False) -> th.Tensor:
return self.actor.forward(observation, deterministic)
MlpPolicy = SACPolicy
register_policy("MlpPolicy", MlpPolicy)

5
torchy_baselines/td3/policies.py
View file

@ -1,12 +1,11 @@
from typing import List, Tuple, Callable, Optional
import torch
import torch as th
import torch.nn as nn
from typing import List, Tuple, Optional
from torchy_baselines.common.distributions import StateDependentNoiseDistribution
from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp, BaseNetwork, \
create_sde_feature_extractor
from torchy_baselines.common.distributions import StateDependentNoiseDistribution
class Actor(BaseNetwork):

14
torchy_baselines/td3/td3.py
View file

@ -1,14 +1,12 @@
from typing import List, Tuple, Type, Union, Callable, Optional, Dict, Any
import torch as th
import torch.nn.functional as F
import numpy as np
from typing import List, Tuple, Type, Union, Callable, Optional, Dict, Any
from torchy_baselines.common.base_class import OffPolicyRLModel
from torchy_baselines.common.buffers import ReplayBuffer
from torchy_baselines.common.type_aliases import ReplayBufferSamples, GymEnv
from torchy_baselines.common.noise import ActionNoise
from torchy_baselines.common.callbacks import BaseCallback
from torchy_baselines.common.noise import ActionNoise
from torchy_baselines.common.type_aliases import ReplayBufferSamples, GymEnv
from torchy_baselines.td3.policies import TD3Policy
@ -138,9 +136,9 @@ class TD3(OffPolicyRLModel):
self.vf_net = self.policy.vf_net
def train_critic(self, gradient_steps: int = 1,
batch_size: int = 100,
replay_data: Optional[ReplayBufferSamples] = None,
tau: float = 0.0) -> None:
batch_size: int = 100,
replay_data: Optional[ReplayBufferSamples] = None,
tau: float = 0.0) -> None:
# Update optimizer learning rate
self._update_learning_rate(self.critic.optimizer)