saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-07-29 20:14:17 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge pull request #71 from Antonin-Raffin/feat/full-policy-save

Browse source 
Independent policy save/load and bug fixes
This commit is contained in:
Raffin, Antonin 2020-04-21 14:56:44 +02:00 committed by GitHub Enterprise
commit 6d2fc328c9
15 changed files with 366 additions and 150 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

55
docs/misc/changelog.rst
View file

@ -3,6 +3,37 @@
Changelog
==========
Pre-Release 0.5.0a1 (WIP)
------------------------------
Breaking Changes:
^^^^^^^^^^^^^^^^^
- Previous loading of policy weights is broken and replace by the new saving/loading for policy
New Features:
^^^^^^^^^^^^^
- Added ``optimizer`` and ``optimizer_kwargs`` to ``policy_kwargs`` in order to easily
customizer optimizers
- Complete independent save/load for policies
Bug Fixes:
^^^^^^^^^^
- Fixed ``reset_num_timesteps`` behavior, so ``env.reset()`` is not called if ``reset_num_timesteps=True``
- Fixed ``squashed_output`` that was not pass to policy constructor for ``SAC`` and ``TD3`` (would result in scaled actions for unscaled action spaces)
Deprecations:
^^^^^^^^^^^^^
Others:
^^^^^^^
- Cleanup rollout return
- Added ``get_device`` util to manage PyTorch devices
Documentation:
^^^^^^^^^^^^^^
Pre-Release 0.4.0 (2020-02-14)
------------------------------
@ -32,10 +63,6 @@ Others:
- Refactored action distributions
Documentation:
^^^^^^^^^^^^^^
Pre-Release 0.3.0 (2020-02-14)
------------------------------
@ -57,9 +84,6 @@ Bug Fixes:
- Fixed colors in ``results_plotter``
- Fix entropy computation (now summed over action dim)
Deprecations:
^^^^^^^^^^^^^
Others:
^^^^^^^
- SAC with SDE now sample only one matrix
@ -106,9 +130,6 @@ Bug Fixes:
- Fix entropy computation for squashed Gaussian (approximate it now)
- Fix seeding when using multiple environments (different seed per env)
Deprecations:
^^^^^^^^^^^^^
Others:
^^^^^^^
- Add type check
@ -125,25 +146,11 @@ Pre-Release 0.1.0 (2020-01-20)
------------------------------
**First Release: base algorithms and state-dependent exploration**
Breaking Changes:
^^^^^^^^^^^^^^^^^
New Features:
^^^^^^^^^^^^^
- Initial release of A2C, CEM-RL, PPO, SAC and TD3, working only with ``Box`` input space
- State-Dependent Exploration (SDE) for A2C, PPO, SAC and TD3
Bug Fixes:
^^^^^^^^^^
Deprecations:
^^^^^^^^^^^^^
Others:
^^^^^^^
Documentation:
^^^^^^^^^^^^^^
Maintainers

24
tests/test_custom_policy.py
View file

@ -1,6 +1,7 @@
import pytest
import torch as th
from torchy_baselines import PPO
from torchy_baselines import A2C, PPO, SAC, TD3
@pytest.mark.parametrize('net_arch', [
@ -11,5 +12,22 @@ from torchy_baselines import PPO
[12, dict(vf=[8], pi=[8, 4])],
[12, dict(pi=[8])],
])
def test_flexible_mlp(net_arch):
_ = PPO('MlpPolicy', 'CartPole-v1', policy_kwargs=dict(net_arch=net_arch), n_steps=100).learn(1000)
@pytest.mark.parametrize('model_class', [A2C, PPO])
def test_flexible_mlp(model_class, net_arch):
_ = model_class('MlpPolicy', 'CartPole-v1', policy_kwargs=dict(net_arch=net_arch), n_steps=100).learn(1000)
@pytest.mark.parametrize('net_arch', [
[4],
[4, 4],
])
@pytest.mark.parametrize('model_class', [SAC, TD3])
def test_custom_offpolicy(model_class, net_arch):
_ = model_class('MlpPolicy', 'Pendulum-v0', policy_kwargs=dict(net_arch=net_arch)).learn(1000)
@pytest.mark.parametrize('model_class', [A2C, PPO, SAC, TD3])
@pytest.mark.parametrize('optimizer_kwargs', [None, dict(weight_decay=0.0)])
def test_custom_optimizer(model_class, optimizer_kwargs):
policy_kwargs = dict(optimizer=th.optim.AdamW, optimizer_kwargs=optimizer_kwargs, net_arch=[32])
_ = model_class('MlpPolicy', 'Pendulum-v0', policy_kwargs=policy_kwargs).learn(1000)

25
tests/test_save_load.py
View file

@ -180,9 +180,11 @@ def test_save_load_policy(model_class):
observations = observations.reshape(10, -1)
policy = model.policy
actor = None
policy_class = policy.__class__
actor, actor_class = None, None
if model_class in [SAC, TD3]:
actor = policy.actor
actor_class = actor.__class__
# Get dictionary of current parameters
params = deepcopy(policy.state_dict())
@ -207,14 +209,17 @@ def test_save_load_policy(model_class):
selected_actions_actor, _ = actor.predict(observations, deterministic=True)
# Save and load policy
policy.save("./logs/policy_weights.pkl")
policy.save("./logs/policy.pkl")
# Save and load actor
if actor is not None:
actor.save("./logs/actor_weights.pkl")
actor.save("./logs/actor.pkl")
policy.load("./logs/policy_weights.pkl")
if actor is not None:
actor.load("./logs/actor_weights.pkl")
del policy, actor
policy = policy_class.load("./logs/policy.pkl")
if actor_class is not None:
actor = actor_class.load("./logs/actor.pkl")
# check if params are still the same after load
new_params = policy.state_dict()
@ -227,12 +232,12 @@ def test_save_load_policy(model_class):
new_selected_actions, _ = policy.predict(observations, deterministic=True)
assert np.allclose(selected_actions, new_selected_actions, 1e-4)
if actor is not None:
if actor_class is not None:
new_selected_actions_actor, _ = actor.predict(observations, deterministic=True)
assert np.allclose(selected_actions_actor, new_selected_actions_actor, 1e-4)
assert np.allclose(selected_actions_actor, new_selected_actions, 1e-4)
# clear file from os
os.remove("./logs/policy_weights.pkl")
if actor is not None:
os.remove("./logs/actor_weights.pkl")
os.remove("./logs/policy.pkl")
if actor_class is not None:
os.remove("./logs/actor.pkl")

14
torchy_baselines/a2c/a2c.py
View file

@ -81,19 +81,15 @@ class A2C(PPO):
seed=seed, _init_setup_model=False)
self.normalize_advantage = normalize_advantage
self.rms_prop_eps = rms_prop_eps
self.use_rms_prop = use_rms_prop
# Override PPO optimizer to match original implementation
if use_rms_prop and 'optimizer' not in self.policy_kwargs:
self.policy_kwargs['optimizer'] = th.optim.RMSprop
self.policy_kwargs['optimizer_kwargs'] = dict(alpha=0.99, eps=rms_prop_eps,
weight_decay=0)
if _init_setup_model:
self._setup_model()
def _setup_model(self) -> None:
super(A2C, self)._setup_model()
if self.use_rms_prop:
self.policy.optimizer = th.optim.RMSprop(self.policy.parameters(),
lr=self.lr_schedule(1), alpha=0.99,
eps=self.rms_prop_eps, weight_decay=0)
def train(self, gradient_steps: int, batch_size: Optional[int] = None) -> None:
# Update optimizer learning rate
self._update_learning_rate(self.policy.optimizer)

70
torchy_baselines/common/base_class.py
View file

@ -13,7 +13,7 @@ import numpy as np
from torchy_baselines.common import logger
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.utils import set_random_seed, get_schedule_fn, update_learning_rate, get_device
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, RolloutReturn, MaybeCallback
@ -71,10 +71,7 @@ class BaseRLModel(ABC):
else:
self.policy_class = policy
if device == 'auto':
device = 'cuda' if th.cuda.is_available() else 'cpu'
self.device = th.device(device)
self.device = get_device(device)
if verbose > 0:
print(f"Using {self.device} device")
@ -93,7 +90,11 @@ class BaseRLModel(ABC):
self.start_time = None
self.policy = None
self.learning_rate = learning_rate
self.lr_schedule = None # type: Optional[Callable]
self.lr_schedule = None # type: Optional[Callable]
self._last_obs = None # type: Optional[np.ndarray]
# When using VecNormalize:
self._last_original_obs = None # type: Optional[np.ndarray]
self._episode_num = 0
# Used for SDE only
self.use_sde = use_sde
self.sde_sample_freq = sde_sample_freq
@ -383,7 +384,7 @@ class BaseRLModel(ABC):
raise ValueError(f"Error: the file {load_path} could not be found")
# set device to cpu if cuda is not available
device = th.device('cuda') if th.cuda.is_available() else th.device('cpu')
device = get_device()
# Open the zip archive and load data
try:
@ -486,7 +487,7 @@ class BaseRLModel(ABC):
n_eval_episodes: int = 5,
log_path: Optional[str] = None,
reset_num_timesteps: bool = True,
) -> Tuple[int, np.ndarray, BaseCallback]:
) -> 'BaseCallback':
"""
Initialize different variables needed for training.
@ -496,7 +497,7 @@ class BaseRLModel(ABC):
:param n_eval_episodes: (int)
:param log_path (Optional[str]): Path to a log folder
:param reset_num_timesteps: (bool) Whether to reset or not the `num_timesteps` attribute
:return: (Tuple[int, np.ndarray, BaseCallback])
:return: (BaseCallback)
"""
self.start_time = time.time()
self.ep_info_buffer = deque(maxlen=100)
@ -505,21 +506,26 @@ class BaseRLModel(ABC):
if self.action_noise is not None:
self.action_noise.reset()
timesteps_since_eval, episode_num = 0, 0
if reset_num_timesteps:
self.num_timesteps = 0
self._episode_num = 0
# Avoid resetting the environment when calling `.learn()` consecutive times
if reset_num_timesteps or self._last_obs is None:
self._last_obs = self.env.reset()
# Retrieve unnormalized observation for saving into the buffer
if self._vec_normalize_env is not None:
self._last_original_obs = self._vec_normalize_env.get_original_obs()
if eval_env is not None and self.seed is not None:
eval_env.seed(self.seed)
eval_env = self._get_eval_env(eval_env)
obs = self.env.reset()
# Create eval callback if needed
callback = self._init_callback(callback, eval_env, eval_freq, n_eval_episodes, log_path)
return episode_num, obs, callback
return callback
def _update_info_buffer(self, infos: List[Dict[str, Any]], dones: Optional[np.ndarray] = None) -> None:
"""
@ -744,8 +750,6 @@ class OffPolicyRLModel(BaseRLModel):
action_noise: Optional[ActionNoise] = None,
learning_starts: int = 0,
replay_buffer: Optional[ReplayBuffer] = None,
obs: Optional[np.ndarray] = None,
episode_num: int = 0,
log_interval: Optional[int] = None) -> RolloutReturn:
"""
Collect rollout using the current policy (and possibly fill the replay buffer)
@ -762,8 +766,6 @@ class OffPolicyRLModel(BaseRLModel):
(and at the beginning and end of the rollout)
:param learning_starts: (int) Number of steps before learning for the warm-up phase.
:param replay_buffer: (ReplayBuffer)
:param obs: (np.ndarray) Last observation from the environment
:param episode_num: (int) Episode index
:param log_interval: (int) Log data every `log_interval` episodes
:return: (RolloutReturn)
"""
@ -773,10 +775,6 @@ class OffPolicyRLModel(BaseRLModel):
assert isinstance(env, VecEnv), "You must pass a VecEnv"
assert env.num_envs == 1, "OffPolicyRLModel only support single environment"
# Retrieve unnormalized observation for saving into the buffer
if self._vec_normalize_env is not None:
obs_ = self._vec_normalize_env.get_original_obs()
self.rollout_data = None
if self.use_sde:
self.actor.reset_noise()
@ -804,7 +802,7 @@ class OffPolicyRLModel(BaseRLModel):
else:
# Note: we assume that the policy uses tanh to scale the action
# We use non-deterministic action in the case of SAC, for TD3, it does not matter
unscaled_action, _ = self.predict(obs, deterministic=False)
unscaled_action, _ = self.predict(self._last_obs, deterministic=False)
# Rescale the action from [low, high] to [-1, 1]
scaled_action = self.policy.scale_action(unscaled_action)
@ -827,7 +825,7 @@ class OffPolicyRLModel(BaseRLModel):
# Only stop training if return value is False, not when it is None.
if callback.on_step() is False:
return RolloutReturn(0.0, total_steps, total_episodes, None, continue_training=False)
return RolloutReturn(0.0, total_steps, total_episodes, continue_training=False)
episode_reward += reward
@ -842,25 +840,23 @@ class OffPolicyRLModel(BaseRLModel):
reward_ = self._vec_normalize_env.get_original_reward()
else:
# Avoid changing the original ones
obs_, new_obs_, reward_ = obs, new_obs, reward
self._last_original_obs, new_obs_, reward_ = self._last_obs, new_obs, reward
replay_buffer.add(obs_, new_obs_, clipped_action, reward_, done)
replay_buffer.add(self._last_original_obs, new_obs_, clipped_action, reward_, done)
if self.rollout_data is not None:
# Assume only one env
self.rollout_data['observations'].append(obs[0].copy())
self.rollout_data['observations'].append(self._last_obs[0].copy())
self.rollout_data['actions'].append(scaled_action[0].copy())
self.rollout_data['rewards'].append(reward[0].copy())
self.rollout_data['dones'].append(done[0].copy())
obs_tensor = th.FloatTensor(obs).to(self.device)
obs_tensor = th.FloatTensor(self._last_obs).to(self.device)
self.rollout_data['values'].append(self.vf_net(obs_tensor)[0].cpu().detach().numpy())
obs = new_obs
# Save the true unnormalized observation
# otherwise obs_ = self._vec_normalize_env.unnormalize_obs(obs)
# is a good approximation
self._last_obs = new_obs
# Save the unnormalized observation
if self._vec_normalize_env is not None:
obs_ = new_obs_
self._last_original_obs = new_obs_
self.num_timesteps += 1
episode_timesteps += 1
@ -870,16 +866,16 @@ class OffPolicyRLModel(BaseRLModel):
if done:
total_episodes += 1
self._episode_num += 1
episode_rewards.append(episode_reward)
total_timesteps.append(episode_timesteps)
if action_noise is not None:
action_noise.reset()
# Display training infos
if self.verbose >= 1 and log_interval is not None and (
episode_num + total_episodes) % log_interval == 0:
if self.verbose >= 1 and log_interval is not None and (self._episode_num) % log_interval == 0:
fps = int(self.num_timesteps / (time.time() - self.start_time))
logger.logkv("episodes", episode_num + total_episodes)
logger.logkv("episodes", self._episode_num)
if len(self.ep_info_buffer) > 0 and len(self.ep_info_buffer[0]) > 0:
logger.logkv('ep_rew_mean', self.safe_mean([ep_info['r'] for ep_info in self.ep_info_buffer]))
logger.logkv('ep_len_mean', self.safe_mean([ep_info['l'] for ep_info in self.ep_info_buffer]))
@ -909,7 +905,7 @@ class OffPolicyRLModel(BaseRLModel):
for step in reversed(range(len(self.rollout_data['rewards']))):
if step == len(self.rollout_data['rewards']) - 1:
next_non_terminal = 1.0 - done[0]
next_value = self.vf_net(th.FloatTensor(obs).to(self.device))[0].detach()
next_value = self.vf_net(th.FloatTensor(self._last_obs).to(self.device))[0].detach()
last_return = self.rollout_data['rewards'][step] + next_non_terminal * next_value
else:
next_non_terminal = 1.0 - self.rollout_data['dones'][step + 1]
@ -919,4 +915,4 @@ class OffPolicyRLModel(BaseRLModel):
callback.on_rollout_end()
return RolloutReturn(mean_reward, total_steps, total_episodes, obs, continue_training)
return RolloutReturn(mean_reward, total_steps, total_episodes, continue_training)

60
torchy_baselines/common/policies.py
View file

@ -1,4 +1,4 @@
from typing import Union, Type, Dict, List, Tuple, Optional
from typing import Union, Type, Dict, List, Tuple, Optional, Any
from itertools import zip_longest
@ -8,6 +8,7 @@ import torch.nn as nn
import numpy as np
from torchy_baselines.common.preprocessing import preprocess_obs
from torchy_baselines.common.utils import get_device, get_schedule_fn
class BasePolicy(nn.Module):
@ -18,7 +19,7 @@ class BasePolicy(nn.Module):
:param action_space: (gym.spaces.Space) The action space of the environment
:param device: (Union[th.device, str]) Device on which the code should run.
:param squash_output: (bool) For continuous actions, whether the output is squashed
or not using a `tanh()` function.
or not using a ``tanh()`` function.
:param features_extractor: (nn.Module) Network to extract features
(a CNN when using images, a nn.Flatten() layer otherwise)
:param normalize_images: (bool) Whether to normalize images or not,
@ -26,17 +27,18 @@ class BasePolicy(nn.Module):
"""
def __init__(self, observation_space: gym.spaces.Space,
action_space: gym.spaces.Space,
device: Union[th.device, str] = 'cpu',
device: Union[th.device, str] = 'auto',
squash_output: bool = False,
features_extractor: Optional[nn.Module] = None,
normalize_images: bool = True):
super(BasePolicy, self).__init__()
self.observation_space = observation_space
self.action_space = action_space
self.device = device
self.device = get_device(device)
self.features_extractor = features_extractor
self.normalize_images = normalize_images
self._squash_output = squash_output
self.optimizer = None # type: Optional[th.optim.Optimizer]
def extract_features(self, obs: th.Tensor) -> th.Tensor:
"""
@ -55,11 +57,19 @@ class BasePolicy(nn.Module):
return self._squash_output
@staticmethod
def init_weights(module: nn.Module, gain: float = 1):
def init_weights(module: nn.Module, gain: float = 1) -> None:
"""
Orthogonal initialization (used in PPO and A2C)
"""
if isinstance(module, nn.Linear):
nn.init.orthogonal_(module.weight, gain=gain)
module.bias.data.fill_(0.0)
@staticmethod
def _dummy_schedule(progress: float) -> float:
""" (float) Useful for pickling policy."""
return 0.0
def forward(self, *_args, **kwargs):
raise NotImplementedError()
@ -189,24 +199,47 @@ class BasePolicy(nn.Module):
raise ValueError("Error: Cannot determine if the observation is vectorized with the space type {}."
.format(observation_space))
def _get_data(self) -> Dict[str, Any]:
"""
Get data that need to be saved in order to re-create the policy.
This corresponds to the arguments of the constructor.
:return: (Dict[str, Any])
"""
return dict(
observation_space=self.observation_space,
action_space=self.action_space,
# Passed to the constructor by child class
# squash_output=self.squash_output,
# features_extractor=self.features_extractor
normalize_images=self.normalize_images,
)
def save(self, path: str) -> None:
"""
Save policy weights to a given location.
NOTE: we don't save policy parameters
Save policy to a given location.
:param path: (str)
"""
th.save(self.state_dict(), path)
th.save({'state_dict': self.state_dict(), 'data': self._get_data()}, path)
def load(self, path: str) -> None:
@classmethod
def load(cls, path: str, device: Union[th.device, str] = 'auto') -> 'BasePolicy':
"""
Load policy weights from path.
NOTE: we don't load policy parameters
Load policy from path.
:param path: (str)
:param device: ( Union[th.device, str]) Device on which the policy should be loaded.
:return: (BasePolicy)
"""
self.load_state_dict(th.load(path))
device = get_device(device)
saved_variables = th.load(path, map_location=device)
# Create policy object
model = cls(**saved_variables['data'])
# Load weights
model.load_state_dict(saved_variables['state_dict'])
model.to(device)
return model
def load_from_vector(self, vector: np.ndarray):
"""
@ -358,8 +391,9 @@ class MlpExtractor(nn.Module):
def __init__(self, feature_dim: int,
net_arch: List[Union[int, Dict[str, List[int]]]],
activation_fn: Type[nn.Module],
device: Union[th.device, str] = 'cpu'):
device: Union[th.device, str] = 'auto'):
super(MlpExtractor, self).__init__()
device = get_device(device)
shared_net, policy_net, value_net = [], [], []
policy_only_layers = [] # Layer sizes of the network that only belongs to the policy network

1
torchy_baselines/common/type_aliases.py
View file

@ -38,5 +38,4 @@ class RolloutReturn(NamedTuple):
episode_reward: float
episode_timesteps: int
n_episodes: int
obs: Optional[np.ndarray]
continue_training: bool

23
torchy_baselines/common/utils.py
View file

@ -84,3 +84,26 @@ def constant_fn(val: float) -> Callable:
return val
return func
def get_device(device: Union[th.device, str] = 'auto') -> th.device:
"""
Retrieve PyTorch device.
It checks that the requested device is available first.
For now, it supports only cpu and cuda.
By default, it tries to use the gpu.
:param device: (Union[str, th.device]) One for 'auto', 'cuda', 'cpu'
:return: (th.device)
"""
# Cuda by default
if device == 'auto':
device = 'cuda'
# Force conversion to th.device
device = th.device(device)
# Cuda not available
if device == th.device('cuda') and not th.cuda.is_available():
return th.device('cpu')
return device

46
torchy_baselines/ppo/policies.py
View file

@ -1,4 +1,4 @@
from typing import Optional, List, Tuple, Callable, Union, Dict, Type
from typing import Optional, List, Tuple, Callable, Union, Dict, Type, Any
from functools import partial
import gym
@ -24,7 +24,6 @@ class PPOPolicy(BasePolicy):
:param net_arch: ([int or dict]) The specification of the policy and value networks.
:param device: (str or th.device) Device on which the code should run.
:param activation_fn: (Type[nn.Module]) Activation function
:param adam_epsilon: (float) Small values to avoid NaN in ADAM optimizer
:param ortho_init: (bool) Whether to use or not orthogonal initialization
:param use_sde: (bool) Whether to use State Dependent Exploration or not
:param log_std_init: (float) Initial value for the log standard deviation
@ -40,15 +39,18 @@ class PPOPolicy(BasePolicy):
this allows to ensure boundaries when using SDE.
:param normalize_images: (bool) Whether to normalize images or not,
dividing by 255.0 (True by default)
:param optimizer: (Type[th.optim.Optimizer]) The optimizer to use,
``th.optim.Adam`` by default
:param optimizer_kwargs: (Optional[Dict[str, Any]]) Additional keyword arguments,
excluding the learning rate, to pass to the optimizer
"""
def __init__(self,
observation_space: gym.spaces.Space,
action_space: gym.spaces.Space,
lr_schedule: Callable,
net_arch: Optional[List[Union[int, Dict[str, List[int]]]]] = None,
device: Union[th.device, str] = 'cpu',
device: Union[th.device, str] = 'auto',
activation_fn: Type[nn.Module] = nn.Tanh,
adam_epsilon: float = 1e-5,
ortho_init: bool = True,
use_sde: bool = False,
log_std_init: float = 0.0,
@ -56,7 +58,9 @@ class PPOPolicy(BasePolicy):
sde_net_arch: Optional[List[int]] = None,
use_expln: bool = False,
squash_output: bool = False,
normalize_images: bool = True):
normalize_images: bool = True,
optimizer: Type[th.optim.Optimizer] = th.optim.Adam,
optimizer_kwargs: Optional[Dict[str, Any]] = None):
super(PPOPolicy, self).__init__(observation_space, action_space, device, squash_output=squash_output)
# Default network architecture, from stable-baselines
@ -65,7 +69,15 @@ class PPOPolicy(BasePolicy):
self.net_arch = net_arch
self.activation_fn = activation_fn
self.adam_epsilon = adam_epsilon
if optimizer_kwargs is None:
optimizer_kwargs = {}
# Small values to avoid NaN in ADAM optimizer
if optimizer == th.optim.Adam:
optimizer_kwargs['eps'] = 1e-5
self.optimizer_class = optimizer
self.optimizer_kwargs = optimizer_kwargs
self.ortho_init = ortho_init
# In the future, feature_extractor will be replaced with a CNN
self.features_extractor = nn.Flatten()
@ -85,12 +97,32 @@ class PPOPolicy(BasePolicy):
self.sde_features_extractor = None
self.sde_net_arch = sde_net_arch
self.use_sde = use_sde
self.dist_kwargs = dist_kwargs
# Action distribution
self.action_dist = make_proba_distribution(action_space, use_sde=use_sde, dist_kwargs=dist_kwargs)
self._build(lr_schedule)
def _get_data(self) -> Dict[str, Any]:
data = super()._get_data()
data.update(dict(
net_arch=self.net_arch,
activation_fn=self.activation_fn,
use_sde=self.use_sde,
log_std_init=self.log_std_init,
squash_output=self.dist_kwargs['squash_output'] if self.dist_kwargs else None,
full_std=self.dist_kwargs['full_std'] if self.dist_kwargs else None,
sde_net_arch=self.dist_kwargs['sde_net_arch'] if self.dist_kwargs else None,
use_expln=self.dist_kwargs['use_expln'] if self.dist_kwargs else None,
lr_schedule=self._dummy_schedule, # dummy lr schedule, not needed for loading policy alone
optimizer=self.optimizer_class,
optimizer_kwargs=self.optimizer_kwargs,
ortho_init=self.ortho_init
))
return data
def reset_noise(self, n_envs: int = 1) -> None:
"""
Sample new weights for the exploration matrix.
@ -142,7 +174,7 @@ class PPOPolicy(BasePolicy):
}[module]
module.apply(partial(self.init_weights, gain=gain))
# Setup optimizer with initial learning rate
self.optimizer = th.optim.Adam(self.parameters(), lr=lr_schedule(1), eps=self.adam_epsilon)
self.optimizer = self.optimizer_class(self.parameters(), lr=lr_schedule(1), **self.optimizer_kwargs)
def forward(self, obs: th.Tensor,
deterministic: bool = False) -> Tuple[th.Tensor, th.Tensor, th.Tensor]:

29
torchy_baselines/ppo/ppo.py
View file

@ -141,12 +141,10 @@ class PPO(BaseRLModel):
env: VecEnv,
callback: BaseCallback,
rollout_buffer: RolloutBuffer,
n_rollout_steps: int = 256,
obs: Optional[np.ndarray] = None) -> Tuple[Optional[np.ndarray], bool]:
n_rollout_steps: int = 256) -> bool:
assert obs is not None, "No previous observation was provided"
assert self._last_obs is not None, "No previous observation was provided"
n_steps = 0
continue_training = True
rollout_buffer.reset()
# Sample new weights for the state dependent exploration
if self.use_sde:
@ -162,7 +160,7 @@ class PPO(BaseRLModel):
with th.no_grad():
# Convert to pytorch tensor
obs_tensor = th.as_tensor(obs).to(self.device)
obs_tensor = th.as_tensor(self._last_obs).to(self.device)
actions, values, log_probs = self.policy.forward(obs_tensor)
actions = actions.cpu().numpy()
@ -171,11 +169,11 @@ class PPO(BaseRLModel):
# Clip the actions to avoid out of bound error
if isinstance(self.action_space, gym.spaces.Box):
clipped_actions = np.clip(actions, self.action_space.low, self.action_space.high)
new_obs, rewards, dones, infos = env.step(clipped_actions)
if callback.on_step() is False:
continue_training = False
return None, continue_training
return False
self._update_info_buffer(infos)
n_steps += 1
@ -184,14 +182,14 @@ class PPO(BaseRLModel):
if isinstance(self.action_space, gym.spaces.Discrete):
# Reshape in case of discrete action
actions = actions.reshape(-1, 1)
rollout_buffer.add(obs, actions, rewards, dones, values, log_probs)
obs = new_obs
rollout_buffer.add(self._last_obs, actions, rewards, dones, values, log_probs)
self._last_obs = new_obs
rollout_buffer.compute_returns_and_advantage(values, dones=dones)
callback.on_rollout_end()
return obs, continue_training
return True
def train(self, n_epochs: int, batch_size: int = 64) -> None:
# Update optimizer learning rate
@ -307,9 +305,9 @@ class PPO(BaseRLModel):
eval_log_path: Optional[str] = None,
reset_num_timesteps: bool = True) -> 'PPO':
episode_num, obs, callback = self._setup_learn(eval_env, callback, eval_freq,
n_eval_episodes, eval_log_path, reset_num_timesteps)
iteration = 0
callback = self._setup_learn(eval_env, callback, eval_freq,
n_eval_episodes, eval_log_path, reset_num_timesteps)
# if self.tensorboard_log is not None and SummaryWriter is not None:
# self.tb_writer = SummaryWriter(log_dir=os.path.join(self.tensorboard_log, tb_log_name))
@ -318,10 +316,9 @@ class PPO(BaseRLModel):
while self.num_timesteps < total_timesteps:
obs, continue_training = self.collect_rollouts(self.env, callback,
self.rollout_buffer,
n_rollout_steps=self.n_steps,
obs=obs)
continue_training = self.collect_rollouts(self.env, callback,
self.rollout_buffer,
n_rollout_steps=self.n_steps)
if continue_training is False:
break

74
torchy_baselines/sac/policies.py
View file

@ -1,4 +1,4 @@
from typing import Optional, List, Tuple, Callable, Union, Type, Dict
from typing import Optional, List, Tuple, Callable, Union, Type, Dict, Any
import gym
import torch as th
@ -53,11 +53,12 @@ class Actor(BasePolicy):
use_expln: bool = False,
clip_mean: float = 2.0,
normalize_images: bool = True,
device: Union[th.device, str] = 'cpu'):
device: Union[th.device, str] = 'auto'):
super(Actor, self).__init__(observation_space, action_space,
features_extractor=features_extractor,
normalize_images=normalize_images,
device=device)
device=device,
squash_output=True)
action_dim = get_action_dim(self.action_space)
@ -65,6 +66,15 @@ class Actor(BasePolicy):
self.latent_pi = nn.Sequential(*latent_pi_net)
self.use_sde = use_sde
self.sde_features_extractor = None
self.sde_net_arch = sde_net_arch
self.net_arch = net_arch
self.features_dim = features_dim
self.activation_fn = activation_fn
self.log_std_init = log_std_init
self.sde_net_arch = sde_net_arch
self.use_expln = use_expln
self.full_std = full_std
self.clip_mean = clip_mean
if self.use_sde:
latent_sde_dim = net_arch[-1]
@ -87,6 +97,23 @@ class Actor(BasePolicy):
self.mu = nn.Linear(net_arch[-1], action_dim)
self.log_std = nn.Linear(net_arch[-1], action_dim)
def _get_data(self) -> Dict[str, Any]:
data = super()._get_data()
data.update(dict(
net_arch=self.net_arch,
features_dim=self.features_dim,
activation_fn=self.activation_fn,
use_sde=self.use_sde,
log_std_init=self.log_std_init,
full_std=self.full_std,
sde_net_arch=self.sde_net_arch,
use_expln=self.use_expln,
features_extractor=self.features_extractor,
clip_mean=self.clip_mean
))
return data
def get_std(self) -> th.Tensor:
"""
Retrieve the standard deviation of the action distribution.
@ -171,7 +198,7 @@ class Critic(BasePolicy):
features_dim: int,
activation_fn: Type[nn.Module] = nn.ReLU,
normalize_images: bool = True,
device: Union[th.device, str] = 'cpu'):
device: Union[th.device, str] = 'auto'):
super(Critic, self).__init__(observation_space, action_space,
features_extractor=features_extractor,
normalize_images=normalize_images,
@ -214,24 +241,36 @@ class SACPolicy(BasePolicy):
:param clip_mean: (float) Clip the mean output when using SDE to avoid numerical instability.
:param normalize_images: (bool) Whether to normalize images or not,
dividing by 255.0 (True by default)
:param optimizer: (Type[th.optim.Optimizer]) The optimizer to use,
``th.optim.Adam`` by default
:param optimizer_kwargs: (Optional[Dict[str, Any]]) Additional keyword arguments,
excluding the learning rate, to pass to the optimizer
"""
def __init__(self, observation_space: gym.spaces.Space,
action_space: gym.spaces.Space,
lr_schedule: Callable,
net_arch: Optional[List[int]] = None,
device: Union[th.device, str] = 'cpu',
device: Union[th.device, str] = 'auto',
activation_fn: Type[nn.Module] = nn.ReLU,
use_sde: bool = False,
log_std_init: float = -3,
sde_net_arch: Optional[List[int]] = None,
use_expln: bool = False,
clip_mean: float = 2.0,
normalize_images: bool = True):
normalize_images: bool = True,
optimizer: Type[th.optim.Optimizer] = th.optim.Adam,
optimizer_kwargs: Optional[Dict[str, Any]] = None):
super(SACPolicy, self).__init__(observation_space, action_space, device, squash_output=True)
if net_arch is None:
net_arch = [256, 256]
if optimizer_kwargs is None:
optimizer_kwargs = {}
self.optimizer_class = optimizer
self.optimizer_kwargs = optimizer_kwargs
# In the future, features_extractor will be replaced with a CNN
self.features_extractor = nn.Flatten()
self.features_dim = get_obs_dim(self.observation_space)
@ -264,12 +303,31 @@ class SACPolicy(BasePolicy):
def _build(self, lr_schedule: Callable) -> None:
self.actor = self.make_actor()
self.actor.optimizer = th.optim.Adam(self.actor.parameters(), lr=lr_schedule(1))
self.actor.optimizer = self.optimizer_class(self.actor.parameters(), lr=lr_schedule(1),
**self.optimizer_kwargs)
self.critic = self.make_critic()
self.critic_target = self.make_critic()
self.critic_target.load_state_dict(self.critic.state_dict())
self.critic.optimizer = th.optim.Adam(self.critic.parameters(), lr=lr_schedule(1))
self.critic.optimizer = self.optimizer_class(self.critic.parameters(), lr=lr_schedule(1),
**self.optimizer_kwargs)
def _get_data(self) -> Dict[str, Any]:
data = super()._get_data()
data.update(dict(
net_arch=self.net_args['net_arch'],
activation_fn=self.net_args['activation_fn'],
use_sde=self.actor_kwargs['use_sde'],
log_std_init=self.actor_kwargs['log_std_init'],
sde_net_arch=self.actor_kwargs['sde_net_arch'],
use_expln=self.actor_kwargs['use_expln'],
clip_mean=self.actor_kwargs['clip_mean'],
lr_schedule=self._dummy_schedule, # dummy lr schedule, not needed for loading policy alone
optimizer=self.optimizer_class,
optimizer_kwargs=self.optimizer_kwargs
))
return data
def make_actor(self) -> Actor:
return Actor(**self.actor_kwargs).to(self.device)

7
torchy_baselines/sac/sac.py
View file

@ -256,8 +256,8 @@ class SAC(OffPolicyRLModel):
eval_log_path: Optional[str] = None,
reset_num_timesteps: bool = True) -> OffPolicyRLModel:
episode_num, obs, callback = self._setup_learn(eval_env, callback, eval_freq,
n_eval_episodes, eval_log_path, reset_num_timesteps)
callback = self._setup_learn(eval_env, callback, eval_freq,
n_eval_episodes, eval_log_path, reset_num_timesteps)
callback.on_training_start(locals(), globals())
while self.num_timesteps < total_timesteps:
@ -266,14 +266,11 @@ class SAC(OffPolicyRLModel):
callback=callback,
learning_starts=self.learning_starts,
replay_buffer=self.replay_buffer,
obs=obs, episode_num=episode_num,
log_interval=log_interval)
if rollout.continue_training is False:
break
obs = rollout.obs
episode_num += rollout.n_episodes
self._update_current_progress(self.num_timesteps, total_timesteps)
if self.num_timesteps > 0 and self.num_timesteps > self.learning_starts:

79
torchy_baselines/td3/policies.py
View file

@ -1,4 +1,4 @@
from typing import Optional, List, Tuple, Callable, Union, Type
from typing import Optional, List, Tuple, Callable, Union, Type, Any, Dict
import gym
import torch as th
@ -52,11 +52,12 @@ class Actor(BasePolicy):
sde_net_arch: Optional[List[int]] = None,
use_expln: bool = False,
normalize_images: bool = True,
device: Union[th.device, str] = 'cpu'):
device: Union[th.device, str] = 'auto'):
super(Actor, self).__init__(observation_space, action_space,
features_extractor=features_extractor,
normalize_images=normalize_images,
device=device)
device=device,
squash_output=not use_sde)
self.latent_pi, self.log_std = None, None
self.weights_dist, self.exploration_mat = None, None
@ -65,6 +66,15 @@ class Actor(BasePolicy):
self.sde_features_extractor = None
self.features_extractor = features_extractor
self.normalize_images = normalize_images
self.net_arch = net_arch
self.features_dim = features_dim
self.activation_fn = activation_fn
self.clip_noise = clip_noise
self.lr_sde = lr_sde
self.log_std_init = log_std_init
self.sde_net_arch = sde_net_arch
self.use_expln = use_expln
self.full_std = full_std
action_dim = get_action_dim(self.action_space)
@ -88,13 +98,30 @@ class Actor(BasePolicy):
log_std_init=log_std_init)
# Squash output
self.mu = nn.Sequential(action_net, nn.Tanh())
self.clip_noise = clip_noise
self.sde_optimizer = th.optim.Adam([self.log_std], lr=lr_sde)
self.reset_noise()
else:
actor_net = create_mlp(features_dim, action_dim, net_arch, activation_fn, squash_output=True)
self.mu = nn.Sequential(*actor_net)
def _get_data(self) -> Dict[str, Any]:
data = super()._get_data()
data.update(dict(
net_arch=self.net_arch,
features_dim=self.features_dim,
activation_fn=self.activation_fn,
use_sde=self.use_sde,
log_std_init=self.log_std_init,
clip_noise=self.clip_noise,
lr_sde=self.lr_sde,
full_std=self.full_std,
sde_net_arch=self.sde_net_arch,
use_expln=self.use_expln,
features_extractor=self.features_extractor
))
return data
def get_std(self) -> th.Tensor:
"""
Retrieve the standard deviation of the action distribution.
@ -179,7 +206,7 @@ class Critic(BasePolicy):
features_dim: int,
activation_fn: Type[nn.Module] = nn.ReLU,
normalize_images: bool = True,
device: Union[th.device, str] = 'cpu'):
device: Union[th.device, str] = 'auto'):
super(Critic, self).__init__(observation_space, action_space,
features_extractor=features_extractor,
normalize_images=normalize_images,
@ -259,12 +286,16 @@ class TD3Policy(BasePolicy):
above zero and prevent it from growing too fast. In practice, ``exp()`` is usually enough.
:param normalize_images: (bool) Whether to normalize images or not,
dividing by 255.0 (True by default)
:param optimizer: (Type[th.optim.Optimizer]) The optimizer to use,
``th.optim.Adam`` by default
:param optimizer_kwargs: (Optional[Dict[str, Any]]) Additional keyword arguments,
excluding the learning rate, to pass to the optimizer
"""
def __init__(self, observation_space: gym.spaces.Space,
action_space: gym.spaces.Space,
lr_schedule: Callable,
net_arch: Optional[List[int]] = None,
device: Union[th.device, str] = 'cpu',
device: Union[th.device, str] = 'auto',
activation_fn: Type[nn.Module] = nn.ReLU,
use_sde: bool = False,
log_std_init: float = -3,
@ -272,13 +303,21 @@ class TD3Policy(BasePolicy):
lr_sde: float = 3e-4,
sde_net_arch: Optional[List[int]] = None,
use_expln: bool = False,
normalize_images: bool = True):
normalize_images: bool = True,
optimizer: Type[th.optim.Optimizer] = th.optim.Adam,
optimizer_kwargs: Optional[Dict[str, Any]] = None):
super(TD3Policy, self).__init__(observation_space, action_space, device, squash_output=True)
# Default network architecture, from the original paper
if net_arch is None:
net_arch = [400, 300]
if optimizer_kwargs is None:
optimizer_kwargs = {}
self.optimizer_class = optimizer
self.optimizer_kwargs = optimizer_kwargs
# In the future, features_extractor will be replaced with a CNN
self.features_extractor = nn.Flatten()
self.features_dim = get_obs_dim(self.observation_space)
@ -318,19 +357,37 @@ class TD3Policy(BasePolicy):
self.actor = self.make_actor()
self.actor_target = self.make_actor()
self.actor_target.load_state_dict(self.actor.state_dict())
self.actor.optimizer = th.optim.Adam(self.actor.parameters(), lr=lr_schedule(1))
self.actor.optimizer = self.optimizer_class(self.actor.parameters(), lr=lr_schedule(1),
**self.optimizer_kwargs)
self.critic = self.make_critic()
self.critic_target = self.make_critic()
self.critic_target.load_state_dict(self.critic.state_dict())
self.critic.optimizer = th.optim.Adam(self.critic.parameters(), lr=lr_schedule(1))
self.critic.optimizer = self.optimizer_class(self.critic.parameters(), lr=lr_schedule(1),
**self.optimizer_kwargs)
if self.use_sde:
self.vf_net = ValueFunction(self.observation_space, self.action_space,
features_extractor=self.features_extractor,
features_dim=self.features_dim)
self.actor.sde_optimizer.add_param_group({'params': self.vf_net.parameters()}) # pytype: disable=attribute-error
def _get_data(self) -> Dict[str, Any]:
data = super()._get_data()
data.update(dict(
net_arch=self.net_args['net_arch'],
activation_fn=self.net_args['activation_fn'],
use_sde=self.actor_kwargs['use_sde'],
log_std_init=self.actor_kwargs['log_std_init'],
clip_noise=self.actor_kwargs['clip_noise'],
lr_sde=self.actor_kwargs['lr_sde'],
sde_net_arch=self.actor_kwargs['sde_net_arch'],
use_expln=self.actor_kwargs['use_expln'],
lr_schedule=self._dummy_schedule, # dummy lr schedule, not needed for loading policy alone
optimizer=self.optimizer_class,
optimizer_kwargs=self.optimizer_kwargs
))
return data
def reset_noise(self) -> None:
return self.actor.reset_noise()

7
torchy_baselines/td3/td3.py
View file

@ -235,8 +235,8 @@ class TD3(OffPolicyRLModel):
eval_log_path: Optional[str] = None,
reset_num_timesteps: bool = True) -> OffPolicyRLModel:
episode_num, obs, callback = self._setup_learn(eval_env, callback, eval_freq,
n_eval_episodes, eval_log_path, reset_num_timesteps)
callback = self._setup_learn(eval_env, callback, eval_freq,
n_eval_episodes, eval_log_path, reset_num_timesteps)
callback.on_training_start(locals(), globals())
@ -247,14 +247,11 @@ class TD3(OffPolicyRLModel):
callback=callback,
learning_starts=self.learning_starts,
replay_buffer=self.replay_buffer,
obs=obs, episode_num=episode_num,
log_interval=log_interval)
if rollout.continue_training is False:
break
obs = rollout.obs
episode_num += rollout.n_episodes
self._update_current_progress(self.num_timesteps, total_timesteps)
if self.num_timesteps > 0 and self.num_timesteps > self.learning_starts:

2
torchy_baselines/version.txt
View file

@ -1 +1 @@
0.4.0
0.5.0a1