mirror of
https://github.com/saymrwulf/stable-baselines3.git
synced 2026-07-29 20:14:17 +00:00
Add proper preprocessing
This commit is contained in:
Antonin RAFFIN
parent
4b2092f55a
commit
ba18258af6
14 changed files with 247 additions and 106 deletions
|
|
@ -23,6 +23,7 @@ Deprecations:
|
|||
Others:
|
||||
^^^^^^^
|
||||
- Refactor handling of observation and action spaces
|
||||
- Refactored features extraction to have proper preprocessing
|
||||
|
||||
Documentation:
|
||||
^^^^^^^^^^^^^^
|
||||
|
|
|
|||
|
|
@ -78,7 +78,7 @@ class BaseRLModel(ABC):
|
|||
if verbose > 0:
|
||||
print(f"Using {self.device} device")
|
||||
|
||||
self.env = None # type: GymEnv
|
||||
self.env = None # type: Optional[GymEnv]
|
||||
# get VecNormalize object if needed
|
||||
self._vec_normalize_env = unwrap_vec_normalize(env)
|
||||
self.verbose = verbose
|
||||
|
|
@ -380,8 +380,10 @@ class BaseRLModel(ABC):
|
|||
:param state: (Optional[np.ndarray]) The last states (can be None, used in recurrent policies)
|
||||
:param mask: (Optional[np.ndarray]) The last masks (can be None, used in recurrent policies)
|
||||
:param deterministic: (bool) Whether or not to return deterministic actions.
|
||||
:return: (Tuple[np.ndarray, Optional[np.ndarray]]) the model's action and the next state (used in recurrent policies)
|
||||
:return: (Tuple[np.ndarray, Optional[np.ndarray]]) the model's action and the next state
|
||||
(used in recurrent policies)
|
||||
"""
|
||||
# TODO: move this block to BasePolicy
|
||||
# if state is None:
|
||||
# state = self.initial_state
|
||||
# if mask is None:
|
||||
|
|
@ -390,9 +392,7 @@ class BaseRLModel(ABC):
|
|||
vectorized_env = self._is_vectorized_observation(observation, self.observation_space)
|
||||
|
||||
observation = observation.reshape((-1,) + self.observation_space.shape)
|
||||
# Convert to float pytorch
|
||||
# TODO: replace with preprocessing
|
||||
observation = th.as_tensor(observation).float().to(self.device)
|
||||
observation = th.as_tensor(observation).to(self.device)
|
||||
with th.no_grad():
|
||||
actions = self.policy.predict(observation, deterministic=deterministic)
|
||||
# Convert to numpy
|
||||
|
|
@ -812,7 +812,7 @@ class OffPolicyRLModel(BaseRLModel):
|
|||
policy_kwargs, verbose,
|
||||
device, support_multi_env, create_eval_env, monitor_wrapper,
|
||||
seed, use_sde, sde_sample_freq)
|
||||
self.buffer_size = buffer_size
|
||||
self.buffer_size = buffer_size
|
||||
self.batch_size = batch_size
|
||||
self.learning_starts = learning_starts
|
||||
self.actor = None
|
||||
|
|
|
|||
|
|
@ -46,7 +46,6 @@ class IdentityEnvBox(IdentityEnv):
|
|||
"""
|
||||
Identity environment for testing purposes
|
||||
|
||||
:param dim: (int) the size of the dimensions you want to learn
|
||||
:param low: (float) the lower bound of the box dim
|
||||
:param high: (float) the upper bound of the box dim
|
||||
:param eps: (float) the epsilon bound for correct value
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
from typing import Union, Type, Dict, List, Tuple
|
||||
from typing import Union, Type, Dict, List, Tuple, Optional
|
||||
|
||||
from itertools import zip_longest
|
||||
|
||||
|
|
@ -7,6 +7,8 @@ import torch as th
|
|||
import torch.nn as nn
|
||||
import numpy as np
|
||||
|
||||
from torchy_baselines.common.preprocessing import preprocess_obs
|
||||
|
||||
|
||||
class BasePolicy(nn.Module):
|
||||
"""
|
||||
|
|
@ -17,17 +19,36 @@ class BasePolicy(nn.Module):
|
|||
: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.
|
||||
: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,
|
||||
dividing by 255.0 (True by default)
|
||||
"""
|
||||
def __init__(self, observation_space: gym.spaces.Space,
|
||||
action_space: gym.spaces.Space,
|
||||
device: Union[th.device, str] = 'cpu',
|
||||
squash_output: bool = False):
|
||||
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.features_extractor = features_extractor
|
||||
self.normalize_images = normalize_images
|
||||
self._squash_output = squash_output
|
||||
|
||||
def extract_features(self, obs: th.Tensor) -> th.Tensor:
|
||||
"""
|
||||
Preprocess the observation if needed and extract features.
|
||||
|
||||
:param obs: (th.Tensor)
|
||||
:return: (th.Tensor)
|
||||
"""
|
||||
assert self.features_extractor is not None, 'No feature extractor was set'
|
||||
preprocessed_obs = preprocess_obs(obs, self.observation_space, normalize_images=self.normalize_images)
|
||||
return self.features_extractor(preprocessed_obs)
|
||||
|
||||
@property
|
||||
def squash_output(self) -> bool:
|
||||
""" (bool) Getter for squash_output."""
|
||||
|
|
@ -45,6 +66,10 @@ class BasePolicy(nn.Module):
|
|||
def predict(self, observation: th.Tensor, deterministic: bool = False) -> th.Tensor:
|
||||
"""
|
||||
Get the action according to the policy for a given observation.
|
||||
|
||||
:param observation: (th.Tensor)
|
||||
:param deterministic: (bool) Whether to use stochastic or deterministic actions
|
||||
:return: (th.Tensor) Taken action according to the policy
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
|
|
@ -118,12 +143,12 @@ def create_mlp(input_dim: int,
|
|||
return modules
|
||||
|
||||
|
||||
def create_sde_feature_extractor(features_dim: int,
|
||||
sde_net_arch: List[int],
|
||||
activation_fn: nn.Module) -> Tuple[nn.Sequential, int]:
|
||||
def create_sde_features_extractor(features_dim: int,
|
||||
sde_net_arch: List[int],
|
||||
activation_fn: nn.Module) -> Tuple[nn.Sequential, int]:
|
||||
"""
|
||||
Create the neural network that will be used to extract features
|
||||
for the SDE.
|
||||
for the SDE exploration function.
|
||||
|
||||
:param features_dim: (int)
|
||||
:param sde_net_arch: ([int])
|
||||
|
|
@ -135,8 +160,8 @@ def create_sde_feature_extractor(features_dim: int,
|
|||
sde_activation = activation_fn if len(sde_net_arch) > 0 else None
|
||||
latent_sde_net = create_mlp(features_dim, -1, sde_net_arch, activation_fn=sde_activation, squash_output=False)
|
||||
latent_sde_dim = sde_net_arch[-1] if len(sde_net_arch) > 0 else features_dim
|
||||
sde_feature_extractor = nn.Sequential(*latent_sde_net)
|
||||
return sde_feature_extractor, latent_sde_dim
|
||||
sde_features_extractor = nn.Sequential(*latent_sde_net)
|
||||
return sde_features_extractor, latent_sde_dim
|
||||
|
||||
|
||||
_policy_registry = dict() # type: Dict[Type[BasePolicy], Dict[str, Type[BasePolicy]]]
|
||||
|
|
|
|||
|
|
@ -34,7 +34,7 @@ def is_image_space(observation_space: spaces.Space) -> bool:
|
|||
|
||||
|
||||
def preprocess_obs(obs: th.Tensor, observation_space: spaces.Space,
|
||||
normalize_image: bool = True) -> th.Tensor:
|
||||
normalize_images: bool = True) -> th.Tensor:
|
||||
"""
|
||||
Preprocess observation to be to a neural network.
|
||||
For images, it normalizes the values by dividing them by 255 (to have values in [0, 1])
|
||||
|
|
@ -42,14 +42,14 @@ def preprocess_obs(obs: th.Tensor, observation_space: spaces.Space,
|
|||
|
||||
:param obs: (th.Tensor) Observation
|
||||
:param observation_space: (spaces.Space)
|
||||
:param normalize_image: (bool) Whether to normalize images or not
|
||||
:param normalize_images: (bool) Whether to normalize images or not
|
||||
(True by default)
|
||||
:return: (th.Tensor)
|
||||
"""
|
||||
if isinstance(observation_space, spaces.Box):
|
||||
if is_image_space(observation_space) and normalize_image:
|
||||
return obs / 255.0
|
||||
return obs
|
||||
if is_image_space(observation_space) and normalize_images:
|
||||
return obs.float() / 255.0
|
||||
return obs.float()
|
||||
elif isinstance(observation_space, spaces.Discrete):
|
||||
# One hot encoding and convert to float to avoid errors
|
||||
return F.one_hot(obs, num_classes=observation_space.n).float()
|
||||
|
|
@ -69,7 +69,7 @@ def get_obs_shape(observation_space: spaces.Space) -> Tuple[int, ...]:
|
|||
return observation_space.shape
|
||||
elif isinstance(observation_space, spaces.Discrete):
|
||||
# Observation is an int
|
||||
return (1,)
|
||||
return 1,
|
||||
else:
|
||||
# TODO: Multidiscrete, Binary, MultiBinary, Tuple, Dict
|
||||
raise NotImplementedError()
|
||||
|
|
|
|||
|
|
@ -5,9 +5,8 @@ used to serialize data (class parameters) of model classes
|
|||
import json
|
||||
import base64
|
||||
import functools
|
||||
from typing import Dict, Any, Optional, Union
|
||||
from typing import Dict, Any, Optional
|
||||
|
||||
import torch as th
|
||||
import cloudpickle
|
||||
import warnings
|
||||
|
||||
|
|
|
|||
|
|
@ -1,7 +1,6 @@
|
|||
"""
|
||||
Common aliases for type hint
|
||||
"""
|
||||
import typing
|
||||
from typing import Union, Dict, Any, NamedTuple, Optional, List, Callable
|
||||
|
||||
import numpy as np
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
# flake8: noqa F401
|
||||
import typing
|
||||
from typing import Optional
|
||||
from typing import Optional, Union
|
||||
from copy import deepcopy
|
||||
|
||||
from torchy_baselines.common.vec_env.base_vec_env import (AlreadySteppingError, NotSteppingError,
|
||||
|
|
@ -15,7 +15,7 @@ if typing.TYPE_CHECKING:
|
|||
from torchy_baselines.common.type_aliases import GymEnv
|
||||
|
||||
|
||||
def unwrap_vec_normalize(env: 'GymEnv') -> Optional[VecNormalize]:
|
||||
def unwrap_vec_normalize(env: Union['GymEnv', VecEnv]) -> Optional[VecNormalize]:
|
||||
"""
|
||||
:param env: (gym.Env)
|
||||
:return: (VecNormalize)
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ import numpy as np
|
|||
|
||||
from torchy_baselines.common.preprocessing import get_obs_dim
|
||||
from torchy_baselines.common.policies import (BasePolicy, register_policy, MlpExtractor,
|
||||
create_sde_feature_extractor)
|
||||
create_sde_features_extractor)
|
||||
from torchy_baselines.common.distributions import (make_proba_distribution, Distribution,
|
||||
DiagGaussianDistribution, CategoricalDistribution,
|
||||
StateDependentNoiseDistribution)
|
||||
|
|
@ -38,6 +38,8 @@ class PPOPolicy(BasePolicy):
|
|||
above zero and prevent it from growing too fast. In practice, ``exp()`` is usually enough.
|
||||
:param squash_output: (bool) Whether to squash the output using a tanh function,
|
||||
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)
|
||||
"""
|
||||
def __init__(self,
|
||||
observation_space: gym.spaces.Space,
|
||||
|
|
@ -53,9 +55,9 @@ class PPOPolicy(BasePolicy):
|
|||
full_std: bool = True,
|
||||
sde_net_arch: Optional[List[int]] = None,
|
||||
use_expln: bool = False,
|
||||
squash_output: bool = False):
|
||||
squash_output: bool = False,
|
||||
normalize_images: bool = True):
|
||||
super(PPOPolicy, self).__init__(observation_space, action_space, device, squash_output=squash_output)
|
||||
self.obs_dim = get_obs_dim(self.observation_space)
|
||||
|
||||
# Default network architecture, from stable-baselines
|
||||
if net_arch is None:
|
||||
|
|
@ -67,7 +69,8 @@ class PPOPolicy(BasePolicy):
|
|||
self.ortho_init = ortho_init
|
||||
# In the future, feature_extractor will be replaced with a CNN
|
||||
self.features_extractor = nn.Flatten()
|
||||
self.features_dim = self.obs_dim
|
||||
self.features_dim = get_obs_dim(self.observation_space)
|
||||
self.normalize_images = normalize_images
|
||||
self.log_std_init = log_std_init
|
||||
dist_kwargs = None
|
||||
# Keyword arguments for SDE distribution
|
||||
|
|
@ -79,7 +82,7 @@ class PPOPolicy(BasePolicy):
|
|||
'learn_features': sde_net_arch is not None
|
||||
}
|
||||
|
||||
self.sde_feature_extractor = None
|
||||
self.sde_features_extractor = None
|
||||
self.sde_net_arch = sde_net_arch
|
||||
self.use_sde = use_sde
|
||||
|
||||
|
|
@ -98,6 +101,12 @@ class PPOPolicy(BasePolicy):
|
|||
self.action_dist.sample_weights(self.log_std, batch_size=n_envs)
|
||||
|
||||
def _build(self, lr_schedule: Callable) -> None:
|
||||
"""
|
||||
Create the networks and the optimizer.
|
||||
|
||||
:param lr_schedule: (Callable) Learning rate schedule
|
||||
lr_schedule(1) is the initial learning rate
|
||||
"""
|
||||
self.mlp_extractor = MlpExtractor(self.features_dim, net_arch=self.net_arch,
|
||||
activation_fn=self.activation_fn, device=self.device)
|
||||
|
||||
|
|
@ -105,9 +114,9 @@ class PPOPolicy(BasePolicy):
|
|||
|
||||
# Separate feature extractor for SDE
|
||||
if self.sde_net_arch is not None:
|
||||
self.sde_feature_extractor, latent_sde_dim = create_sde_feature_extractor(self.features_dim,
|
||||
self.sde_net_arch,
|
||||
self.activation_fn)
|
||||
self.sde_features_extractor, latent_sde_dim = create_sde_features_extractor(self.features_dim,
|
||||
self.sde_net_arch,
|
||||
self.activation_fn)
|
||||
|
||||
if isinstance(self.action_dist, DiagGaussianDistribution):
|
||||
self.action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=latent_dim_pi,
|
||||
|
|
@ -132,12 +141,20 @@ class PPOPolicy(BasePolicy):
|
|||
self.value_net: 1
|
||||
}[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)
|
||||
|
||||
def forward(self, obs: th.Tensor, deterministic: bool = False) -> Tuple[th.Tensor, th.Tensor, th.Tensor]:
|
||||
if not isinstance(obs, th.Tensor):
|
||||
obs = th.FloatTensor(obs).to(self.device)
|
||||
def forward(self, obs: th.Tensor,
|
||||
deterministic: bool = False) -> Tuple[th.Tensor, th.Tensor, th.Tensor]:
|
||||
"""
|
||||
Forward pass in all the networks (actor and critic)
|
||||
|
||||
:param obs: (th.Tensor) Observation
|
||||
:param deterministic: (bool) Whether to sample or use deterministic actions
|
||||
:return: (Tuple[th.Tensor, th.Tensor, th.Tensor]) action, value and log probability of the action
|
||||
"""
|
||||
latent_pi, latent_vf, latent_sde = self._get_latent(obs)
|
||||
# Evaluate the values for the given observations
|
||||
value = self.value_net(latent_vf)
|
||||
action, action_distribution = self._get_action_dist_from_latent(latent_pi, latent_sde=latent_sde,
|
||||
deterministic=deterministic)
|
||||
|
|
@ -145,17 +162,35 @@ class PPOPolicy(BasePolicy):
|
|||
return action, value, log_prob
|
||||
|
||||
def _get_latent(self, obs: th.Tensor) -> Tuple[th.Tensor, th.Tensor, th.Tensor]:
|
||||
features = self.features_extractor(obs)
|
||||
"""
|
||||
Get the latent code (i.e., activations of the last layer of each network)
|
||||
for the different networks.
|
||||
|
||||
:param obs: (th.Tensor) Observation
|
||||
:return: (Tuple[th.Tensor, th.Tensor, th.Tensor]) Latent codes
|
||||
for the actor, the value function and for SDE function
|
||||
"""
|
||||
# Preprocess the observation if needed
|
||||
features = self.extract_features(obs)
|
||||
latent_pi, latent_vf = self.mlp_extractor(features)
|
||||
# Features for sde
|
||||
latent_sde = latent_pi
|
||||
if self.sde_feature_extractor is not None:
|
||||
latent_sde = self.sde_feature_extractor(features)
|
||||
if self.sde_features_extractor is not None:
|
||||
latent_sde = self.sde_features_extractor(features)
|
||||
return latent_pi, latent_vf, latent_sde
|
||||
|
||||
def _get_action_dist_from_latent(self, latent_pi: th.Tensor,
|
||||
latent_sde: Optional[th.Tensor] = None,
|
||||
deterministic: bool = False) -> Tuple[th.Tensor, Distribution]:
|
||||
"""
|
||||
Retrieve action and associated action distribution
|
||||
given the latent codes.
|
||||
|
||||
:param latent_pi: (th.Tensor) Latent code for the actor
|
||||
:param latent_sde: (Optional[th.Tensor]) Latent code for the SDE exploration function
|
||||
:param deterministic: (bool) Whether to sample or use deterministic actions
|
||||
:return: (Tuple[th.Tensor, Distribution]) Action and action distribution
|
||||
"""
|
||||
mean_actions = self.action_net(latent_pi)
|
||||
|
||||
if isinstance(self.action_dist, DiagGaussianDistribution):
|
||||
|
|
@ -172,6 +207,13 @@ class PPOPolicy(BasePolicy):
|
|||
raise ValueError('Invalid action distribution')
|
||||
|
||||
def predict(self, observation: th.Tensor, deterministic: bool = False) -> th.Tensor:
|
||||
"""
|
||||
Get the action according to the policy for a given observation.
|
||||
|
||||
:param observation: (th.Tensor)
|
||||
:param deterministic: (bool) Whether to use stochastic or deterministic actions
|
||||
:return: (th.Tensor) Taken action according to the policy
|
||||
"""
|
||||
latent_pi, _, latent_sde = self._get_latent(observation)
|
||||
action, _ = self._get_action_dist_from_latent(latent_pi, latent_sde, deterministic=deterministic)
|
||||
return action
|
||||
|
|
|
|||
|
|
@ -1,4 +1,3 @@
|
|||
import os
|
||||
import time
|
||||
from typing import List, Tuple, Type, Union, Callable, Optional, Dict, Any
|
||||
|
||||
|
|
@ -8,10 +7,11 @@ import torch as th
|
|||
import torch.nn.functional as F
|
||||
|
||||
# Check if tensorboard is available for pytorch
|
||||
try:
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
except ImportError:
|
||||
SummaryWriter = None
|
||||
# TODO: finish tensorboard integration
|
||||
# try:
|
||||
# from torch.utils.tensorboard import SummaryWriter
|
||||
# except ImportError:
|
||||
# SummaryWriter = None
|
||||
import numpy as np
|
||||
|
||||
from torchy_baselines.common import logger
|
||||
|
|
@ -144,6 +144,7 @@ class PPO(BaseRLModel):
|
|||
n_rollout_steps: int = 256,
|
||||
obs: Optional[np.ndarray] = None) -> Tuple[Optional[np.ndarray], bool]:
|
||||
|
||||
assert obs is not None, "No previous observation was provided"
|
||||
n_steps = 0
|
||||
continue_training = True
|
||||
rollout_buffer.reset()
|
||||
|
|
@ -160,7 +161,10 @@ class PPO(BaseRLModel):
|
|||
self.policy.reset_noise(env.num_envs)
|
||||
|
||||
with th.no_grad():
|
||||
actions, values, log_probs = self.policy.forward(obs)
|
||||
# Convert to pytorch tensor
|
||||
obs_tensor = obs.reshape((-1,) + self.observation_space.shape)
|
||||
obs_tensor = th.as_tensor(obs_tensor).to(self.device)
|
||||
actions, values, log_probs = self.policy.forward(obs_tensor)
|
||||
actions = actions.cpu().numpy()
|
||||
|
||||
# Rescale and perform action
|
||||
|
|
@ -308,8 +312,8 @@ class PPO(BaseRLModel):
|
|||
n_eval_episodes, eval_log_path, reset_num_timesteps)
|
||||
iteration = 0
|
||||
|
||||
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))
|
||||
# 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))
|
||||
|
||||
callback.on_training_start(locals(), globals())
|
||||
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ import torch.nn as nn
|
|||
|
||||
from torchy_baselines.common.preprocessing import get_action_dim, get_obs_dim
|
||||
from torchy_baselines.common.policies import (BasePolicy, register_policy, create_mlp,
|
||||
create_sde_feature_extractor)
|
||||
create_sde_features_extractor)
|
||||
from torchy_baselines.common.distributions import SquashedDiagGaussianDistribution, StateDependentNoiseDistribution
|
||||
|
||||
# CAP the standard deviation of the actor
|
||||
|
|
@ -21,6 +21,9 @@ class Actor(BasePolicy):
|
|||
:param observation_space: (gym.spaces.Space) Obervation space
|
||||
:param action_space: (gym.spaces.Space) Action space
|
||||
:param net_arch: ([int]) Network architecture
|
||||
:param features_extractor: (nn.Module) Network to extract features
|
||||
(a CNN when using images, a nn.Flatten() layer otherwise)
|
||||
:param features_dim: (int) Number of features
|
||||
:param activation_fn: (nn.Module) Activation function
|
||||
:param use_sde: (bool) Whether to use State Dependent Exploration or not
|
||||
:param log_std_init: (float) Initial value for the log standard deviation
|
||||
|
|
@ -33,33 +36,39 @@ class Actor(BasePolicy):
|
|||
a positive standard deviation (cf paper). It allows to keep variance
|
||||
above zero and prevent it from growing too fast. In practice, ``exp()`` is usually enough.
|
||||
: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)
|
||||
"""
|
||||
def __init__(self, observation_space: gym.spaces.Space,
|
||||
action_space: gym.spaces.Space,
|
||||
net_arch: List[int],
|
||||
features_extractor: nn.Module,
|
||||
features_dim: int,
|
||||
activation_fn: nn.Module = nn.ReLU,
|
||||
use_sde: bool = False,
|
||||
log_std_init: float = -3,
|
||||
full_std: bool = True,
|
||||
sde_net_arch: Optional[List[int]] = None,
|
||||
use_expln: bool = False,
|
||||
clip_mean: float = 2.0):
|
||||
super(Actor, self).__init__(observation_space, action_space)
|
||||
clip_mean: float = 2.0,
|
||||
normalize_images: bool = True):
|
||||
super(Actor, self).__init__(observation_space, action_space,
|
||||
features_extractor=features_extractor,
|
||||
normalize_images=normalize_images)
|
||||
|
||||
obs_dim = get_obs_dim(self.observation_space)
|
||||
action_dim = get_action_dim(self.action_space)
|
||||
|
||||
latent_pi_net = create_mlp(obs_dim, -1, net_arch, activation_fn)
|
||||
latent_pi_net = create_mlp(features_dim, -1, net_arch, activation_fn)
|
||||
self.latent_pi = nn.Sequential(*latent_pi_net)
|
||||
self.use_sde = use_sde
|
||||
self.sde_feature_extractor = None
|
||||
self.sde_features_extractor = None
|
||||
|
||||
if self.use_sde:
|
||||
latent_sde_dim = net_arch[-1]
|
||||
# Separate feature extractor for SDE
|
||||
if sde_net_arch is not None:
|
||||
self.sde_feature_extractor, latent_sde_dim = create_sde_feature_extractor(obs_dim, sde_net_arch,
|
||||
activation_fn)
|
||||
self.sde_features_extractor, latent_sde_dim = create_sde_features_extractor(features_dim, sde_net_arch,
|
||||
activation_fn)
|
||||
|
||||
self.action_dist = StateDependentNoiseDistribution(action_dim, full_std=full_std, use_expln=use_expln,
|
||||
learn_features=True, squash_output=True)
|
||||
|
|
@ -85,7 +94,8 @@ class Actor(BasePolicy):
|
|||
|
||||
:return: (th.Tensor)
|
||||
"""
|
||||
assert isinstance(self.action_dist, StateDependentNoiseDistribution), 'get_std() is only available when using SDE'
|
||||
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: int = 1) -> None:
|
||||
|
|
@ -94,14 +104,14 @@ class Actor(BasePolicy):
|
|||
|
||||
:param batch_size: (int)
|
||||
"""
|
||||
assert isinstance(self.action_dist, StateDependentNoiseDistribution), 'reset_noise() is only available when using SDE'
|
||||
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: th.Tensor) -> Tuple[th.Tensor, th.Tensor]:
|
||||
latent_pi = self.latent_pi(obs)
|
||||
|
||||
latent_sde = self.sde_feature_extractor(obs) if self.sde_feature_extractor is not None else latent_pi
|
||||
|
||||
features = self.extract_features(obs)
|
||||
latent_pi = self.latent_pi(features)
|
||||
latent_sde = self.sde_features_extractor(features) if self.sde_features_extractor is not None else latent_pi
|
||||
return latent_pi, latent_sde
|
||||
|
||||
def get_action_dist_params(self, obs: th.Tensor) -> Tuple[th.Tensor, th.Tensor, th.Tensor]:
|
||||
|
|
@ -138,27 +148,37 @@ class Critic(BasePolicy):
|
|||
:param observation_space: (gym.spaces.Space) Obervation space
|
||||
:param action_space: (gym.spaces.Space) Action space
|
||||
:param net_arch: ([int]) Network architecture
|
||||
:param features_extractor: (nn.Module) Network to extract features
|
||||
(a CNN when using images, a nn.Flatten() layer otherwise)
|
||||
:param features_dim: (int) Number of features
|
||||
:param activation_fn: (nn.Module) Activation function
|
||||
:param normalize_images: (bool) Whether to normalize images or not,
|
||||
dividing by 255.0 (True by default)
|
||||
"""
|
||||
def __init__(self, observation_space: gym.spaces.Space,
|
||||
action_space: gym.spaces.Space,
|
||||
net_arch: List[int],
|
||||
activation_fn: nn.Module = nn.ReLU):
|
||||
super(Critic, self).__init__(observation_space, action_space)
|
||||
features_extractor: nn.Module,
|
||||
features_dim: int,
|
||||
activation_fn: nn.Module = nn.ReLU,
|
||||
normalize_images: bool = True):
|
||||
super(Critic, self).__init__(observation_space, action_space,
|
||||
features_extractor=features_extractor,
|
||||
normalize_images=normalize_images)
|
||||
|
||||
obs_dim = get_obs_dim(self.observation_space)
|
||||
action_dim = get_action_dim(self.action_space)
|
||||
|
||||
q1_net = create_mlp(obs_dim + action_dim, 1, net_arch, activation_fn)
|
||||
q1_net = create_mlp(features_dim + action_dim, 1, net_arch, activation_fn)
|
||||
self.q1_net = nn.Sequential(*q1_net)
|
||||
|
||||
q2_net = create_mlp(obs_dim + action_dim, 1, net_arch, activation_fn)
|
||||
q2_net = create_mlp(features_dim + action_dim, 1, net_arch, activation_fn)
|
||||
self.q2_net = nn.Sequential(*q2_net)
|
||||
|
||||
self.q_networks = [self.q1_net, self.q2_net]
|
||||
|
||||
def forward(self, obs: th.Tensor, action: th.Tensor) -> List[th.Tensor]:
|
||||
qvalue_input = th.cat([obs, action], dim=1)
|
||||
features = self.extract_features(obs)
|
||||
qvalue_input = th.cat([features, action], dim=1)
|
||||
return [q_net(qvalue_input) for q_net in self.q_networks]
|
||||
|
||||
|
||||
|
|
@ -181,6 +201,8 @@ class SACPolicy(BasePolicy):
|
|||
a positive standard deviation (cf paper). It allows to keep variance
|
||||
above zero and prevent it from growing too fast. In practice, ``exp()`` is usually enough.
|
||||
: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)
|
||||
"""
|
||||
def __init__(self, observation_space: gym.spaces.Space,
|
||||
action_space: gym.spaces.Space,
|
||||
|
|
@ -192,19 +214,27 @@ class SACPolicy(BasePolicy):
|
|||
log_std_init: float = -3,
|
||||
sde_net_arch: Optional[List[int]] = None,
|
||||
use_expln: bool = False,
|
||||
clip_mean: float = 2.0):
|
||||
clip_mean: float = 2.0,
|
||||
normalize_images: bool = True):
|
||||
super(SACPolicy, self).__init__(observation_space, action_space, device, squash_output=True)
|
||||
|
||||
if net_arch is None:
|
||||
net_arch = [256, 256]
|
||||
|
||||
# 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)
|
||||
|
||||
self.net_arch = net_arch
|
||||
self.activation_fn = activation_fn
|
||||
self.net_args = {
|
||||
'observation_space': self.observation_space,
|
||||
'action_space': self.action_space,
|
||||
'features_extractor': self.features_extractor,
|
||||
'features_dim': self.features_dim,
|
||||
'net_arch': self.net_arch,
|
||||
'activation_fn': self.activation_fn
|
||||
'activation_fn': self.activation_fn,
|
||||
'normalize_images': normalize_images
|
||||
}
|
||||
self.actor_kwargs = self.net_args.copy()
|
||||
sde_kwargs = {
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ import torch.nn as nn
|
|||
|
||||
from torchy_baselines.common.preprocessing import get_action_dim, get_obs_dim
|
||||
from torchy_baselines.common.policies import (BasePolicy, register_policy, create_mlp,
|
||||
create_sde_feature_extractor)
|
||||
create_sde_features_extractor)
|
||||
from torchy_baselines.common.distributions import StateDependentNoiseDistribution, Distribution
|
||||
|
||||
|
||||
|
|
@ -17,6 +17,9 @@ class Actor(BasePolicy):
|
|||
:param observation_space: (gym.spaces.Space) Obervation space
|
||||
:param action_space: (gym.spaces.Space) Action space
|
||||
:param net_arch: ([int]) Network architecture
|
||||
:param features_extractor: (nn.Module) Network to extract features
|
||||
(a CNN when using images, a nn.Flatten() layer otherwise)
|
||||
:param features_dim: (int) Number of features
|
||||
:param activation_fn: (nn.Module) Activation function
|
||||
:param use_sde: (bool) Whether to use State Dependent Exploration or not
|
||||
:param log_std_init: (float) Initial value for the log standard deviation
|
||||
|
|
@ -30,11 +33,15 @@ class Actor(BasePolicy):
|
|||
:param use_expln: (bool) Use ``expln()`` function instead of ``exp()`` when using SDE to ensure
|
||||
a positive standard deviation (cf paper). It allows to keep variance
|
||||
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)
|
||||
"""
|
||||
def __init__(self,
|
||||
observation_space: gym.spaces.Space,
|
||||
action_space: gym.spaces.Space,
|
||||
net_arch: List[int],
|
||||
features_extractor: nn.Module,
|
||||
features_dim: int,
|
||||
activation_fn: nn.Module = nn.ReLU,
|
||||
use_sde: bool = False,
|
||||
log_std_init: float = -3,
|
||||
|
|
@ -42,32 +49,37 @@ class Actor(BasePolicy):
|
|||
lr_sde: float = 3e-4,
|
||||
full_std: bool = False,
|
||||
sde_net_arch: Optional[List[int]] = None,
|
||||
use_expln: bool = False):
|
||||
super(Actor, self).__init__(observation_space, action_space)
|
||||
use_expln: bool = False,
|
||||
normalize_images: bool = True):
|
||||
super(Actor, self).__init__(observation_space, action_space,
|
||||
features_extractor=features_extractor,
|
||||
normalize_images=normalize_images)
|
||||
|
||||
self.latent_pi, self.log_std = None, None
|
||||
self.weights_dist, self.exploration_mat = None, None
|
||||
self.use_sde, self.sde_optimizer = use_sde, None
|
||||
self.full_std = full_std
|
||||
self.sde_feature_extractor = None
|
||||
self.sde_features_extractor = None
|
||||
self.features_extractor = features_extractor
|
||||
self.normalize_images = normalize_images
|
||||
|
||||
obs_dim = get_obs_dim(self.observation_space)
|
||||
action_dim = get_action_dim(self.action_space)
|
||||
|
||||
if use_sde:
|
||||
latent_pi_net = create_mlp(obs_dim, -1, net_arch, activation_fn, squash_output=False)
|
||||
latent_pi_net = create_mlp(features_dim, -1, net_arch, activation_fn, squash_output=False)
|
||||
self.latent_pi = nn.Sequential(*latent_pi_net)
|
||||
latent_sde_dim = net_arch[-1]
|
||||
learn_features = sde_net_arch is not None
|
||||
|
||||
# Separate feature extractor for SDE
|
||||
if sde_net_arch is not None:
|
||||
self.sde_feature_extractor, latent_sde_dim = create_sde_feature_extractor(obs_dim, sde_net_arch,
|
||||
activation_fn)
|
||||
self.sde_features_extractor, latent_sde_dim = create_sde_features_extractor(features_dim, sde_net_arch,
|
||||
activation_fn)
|
||||
|
||||
# Create state dependent noise matrix (SDE)
|
||||
self.action_dist = StateDependentNoiseDistribution(action_dim, full_std=full_std, use_expln=use_expln,
|
||||
squash_output=False, learn_features=learn_features)
|
||||
|
||||
action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=net_arch[-1],
|
||||
latent_sde_dim=latent_sde_dim,
|
||||
log_std_init=log_std_init)
|
||||
|
|
@ -77,7 +89,7 @@ class Actor(BasePolicy):
|
|||
self.sde_optimizer = th.optim.Adam([self.log_std], lr=lr_sde)
|
||||
self.reset_noise()
|
||||
else:
|
||||
actor_net = create_mlp(obs_dim, action_dim, net_arch, activation_fn, squash_output=True)
|
||||
actor_net = create_mlp(features_dim, action_dim, net_arch, activation_fn, squash_output=True)
|
||||
self.mu = nn.Sequential(*actor_net)
|
||||
|
||||
def get_std(self) -> th.Tensor:
|
||||
|
|
@ -98,12 +110,9 @@ class Actor(BasePolicy):
|
|||
return self.action_dist.proba_distribution(mean_actions, self.log_std, latent_sde)
|
||||
|
||||
def _get_latent(self, obs: th.Tensor) -> Tuple[th.Tensor, th.Tensor]:
|
||||
latent_pi = self.latent_pi(obs)
|
||||
|
||||
if self.sde_feature_extractor is not None:
|
||||
latent_sde = self.sde_feature_extractor(obs)
|
||||
else:
|
||||
latent_sde = latent_pi
|
||||
features = self.extract_features(obs)
|
||||
latent_pi = self.latent_pi(features)
|
||||
latent_sde = self.sde_features_extractor(features) if self.sde_features_extractor is not None else latent_pi
|
||||
return latent_pi, latent_sde
|
||||
|
||||
def evaluate_actions(self, obs: th.Tensor, action: th.Tensor) -> Tuple[th.Tensor, th.Tensor]:
|
||||
|
|
@ -144,7 +153,8 @@ class Actor(BasePolicy):
|
|||
# action, _ = self._get_action_dist_from_latent(latent_pi)
|
||||
# return action
|
||||
else:
|
||||
return self.mu(obs)
|
||||
features = self.extract_features(obs)
|
||||
return self.mu(features)
|
||||
|
||||
|
||||
class Critic(BasePolicy):
|
||||
|
|
@ -155,29 +165,40 @@ class Critic(BasePolicy):
|
|||
:param observation_space: (gym.spaces.Space) Obervation space
|
||||
:param action_space: (gym.spaces.Space) Action space
|
||||
:param net_arch: ([int]) Network architecture
|
||||
:param features_extractor: (nn.Module) Network to extract features
|
||||
(a CNN when using images, a nn.Flatten() layer otherwise)
|
||||
:param features_dim: (int) Number of features
|
||||
:param activation_fn: (nn.Module) Activation function
|
||||
:param normalize_images: (bool) Whether to normalize images or not,
|
||||
dividing by 255.0 (True by default)
|
||||
"""
|
||||
def __init__(self, observation_space: gym.spaces.Space,
|
||||
action_space: gym.spaces.Space,
|
||||
net_arch: List[int],
|
||||
activation_fn: nn.Module = nn.ReLU):
|
||||
super(Critic, self).__init__(observation_space, action_space)
|
||||
features_extractor: nn.Module,
|
||||
features_dim: int,
|
||||
activation_fn: nn.Module = nn.ReLU,
|
||||
normalize_images: bool = True):
|
||||
super(Critic, self).__init__(observation_space, action_space,
|
||||
features_extractor=features_extractor,
|
||||
normalize_images=normalize_images)
|
||||
|
||||
obs_dim = get_obs_dim(self.observation_space)
|
||||
action_dim = get_action_dim(self.action_space)
|
||||
|
||||
q1_net = create_mlp(obs_dim + action_dim, 1, net_arch, activation_fn)
|
||||
q1_net = create_mlp(features_dim + action_dim, 1, net_arch, activation_fn)
|
||||
self.q1_net = nn.Sequential(*q1_net)
|
||||
|
||||
q2_net = create_mlp(obs_dim + action_dim, 1, net_arch, activation_fn)
|
||||
q2_net = create_mlp(features_dim + action_dim, 1, net_arch, activation_fn)
|
||||
self.q2_net = nn.Sequential(*q2_net)
|
||||
|
||||
def forward(self, obs: th.Tensor, action: th.Tensor) -> Tuple[th.Tensor, th.Tensor]:
|
||||
qvalue_input = th.cat([obs, action], dim=1)
|
||||
features = self.extract_features(obs)
|
||||
qvalue_input = th.cat([features, action], dim=1)
|
||||
return self.q1_net(qvalue_input), self.q2_net(qvalue_input)
|
||||
|
||||
def q1_forward(self, obs: th.Tensor, action: th.Tensor) -> th.Tensor:
|
||||
return self.q1_net(th.cat([obs, action], dim=1))
|
||||
features = self.extract_features(obs)
|
||||
return self.q1_net(th.cat([features, action], dim=1))
|
||||
|
||||
|
||||
class ValueFunction(BasePolicy):
|
||||
|
|
@ -186,24 +207,34 @@ class ValueFunction(BasePolicy):
|
|||
|
||||
:param observation_space: (gym.spaces.Space) Obervation space
|
||||
:param action_space: (gym.spaces.Space) Action space
|
||||
:param features_extractor: (nn.Module) Network to extract features
|
||||
(a CNN when using images, a nn.Flatten() layer otherwise)
|
||||
:param features_dim: (int) Number of features
|
||||
:param net_arch: (Optional[List[int]]) Network architecture
|
||||
:param activation_fn: (nn.Module) Activation function
|
||||
:param normalize_images: (bool) Whether to normalize images or not,
|
||||
dividing by 255.0 (True by default)
|
||||
"""
|
||||
def __init__(self, observation_space: gym.spaces.Space,
|
||||
action_space: gym.spaces.Space,
|
||||
features_extractor: nn.Module,
|
||||
features_dim: int,
|
||||
net_arch: Optional[List[int]] = None,
|
||||
activation_fn: nn.Module = nn.Tanh):
|
||||
super(ValueFunction, self).__init__(observation_space, action_space)
|
||||
activation_fn: nn.Module = nn.Tanh,
|
||||
normalize_images: bool = True):
|
||||
super(ValueFunction, self).__init__(observation_space, action_space,
|
||||
features_extractor=features_extractor,
|
||||
normalize_images=normalize_images)
|
||||
|
||||
obs_dim = get_obs_dim(self.observation_space)
|
||||
if net_arch is None:
|
||||
net_arch = [64, 64]
|
||||
|
||||
vf_net = create_mlp(obs_dim, 1, net_arch, activation_fn)
|
||||
vf_net = create_mlp(features_dim, 1, net_arch, activation_fn)
|
||||
self.vf_net = nn.Sequential(*vf_net)
|
||||
|
||||
def forward(self, obs: th.Tensor) -> th.Tensor:
|
||||
return self.vf_net(obs)
|
||||
features = self.extract_features(obs)
|
||||
return self.vf_net(features)
|
||||
|
||||
|
||||
class TD3Policy(BasePolicy):
|
||||
|
|
@ -224,6 +255,8 @@ class TD3Policy(BasePolicy):
|
|||
:param use_expln: (bool) Use ``expln()`` function instead of ``exp()`` when using SDE to ensure
|
||||
a positive standard deviation (cf paper). It allows to keep variance
|
||||
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)
|
||||
"""
|
||||
def __init__(self, observation_space: gym.spaces.Space,
|
||||
action_space: gym.spaces.Space,
|
||||
|
|
@ -236,20 +269,28 @@ class TD3Policy(BasePolicy):
|
|||
clip_noise: Optional[float] = None,
|
||||
lr_sde: float = 3e-4,
|
||||
sde_net_arch: Optional[List[int]] = None,
|
||||
use_expln: bool = False):
|
||||
use_expln: bool = False,
|
||||
normalize_images: bool = True):
|
||||
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]
|
||||
|
||||
# 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)
|
||||
|
||||
self.net_arch = net_arch
|
||||
self.activation_fn = activation_fn
|
||||
self.net_args = {
|
||||
'observation_space': self.observation_space,
|
||||
'action_space': self.action_space,
|
||||
'features_extractor': self.features_extractor,
|
||||
'features_dim': self.features_dim,
|
||||
'net_arch': self.net_arch,
|
||||
'activation_fn': self.activation_fn
|
||||
'activation_fn': self.activation_fn,
|
||||
'normalize_images': normalize_images
|
||||
}
|
||||
self.actor_kwargs = self.net_args.copy()
|
||||
sde_kwargs = {
|
||||
|
|
@ -282,7 +323,9 @@ class TD3Policy(BasePolicy):
|
|||
self.critic.optimizer = th.optim.Adam(self.critic.parameters(), lr=lr_schedule(1))
|
||||
|
||||
if self.use_sde:
|
||||
self.vf_net = ValueFunction(self.observation_space, self.action_space)
|
||||
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 reset_noise(self) -> None:
|
||||
|
|
|
|||
|
|
@ -4,9 +4,8 @@ from typing import List, Tuple, Type, Union, Callable, Optional, Dict, Any
|
|||
|
||||
from torchy_baselines.common import logger
|
||||
from torchy_baselines.common.base_class import OffPolicyRLModel
|
||||
from torchy_baselines.common.buffers import ReplayBuffer
|
||||
from torchy_baselines.common.noise import ActionNoise
|
||||
from torchy_baselines.common.type_aliases import ReplayBufferSamples, GymEnv, MaybeCallback
|
||||
from torchy_baselines.common.type_aliases import GymEnv, MaybeCallback
|
||||
from torchy_baselines.td3.policies import TD3Policy
|
||||
|
||||
|
||||
|
|
@ -161,7 +160,8 @@ class TD3(OffPolicyRLModel):
|
|||
# Delayed policy updates
|
||||
if gradient_step % policy_delay == 0:
|
||||
# Compute actor loss
|
||||
actor_loss = -self.critic.q1_forward(replay_data.observations, self.actor(replay_data.observations)).mean()
|
||||
actor_loss = -self.critic.q1_forward(replay_data.observations,
|
||||
self.actor(replay_data.observations)).mean()
|
||||
|
||||
# Optimize the actor
|
||||
self.actor.optimizer.zero_grad()
|
||||
|
|
@ -178,7 +178,6 @@ class TD3(OffPolicyRLModel):
|
|||
self._n_updates += gradient_steps
|
||||
logger.logkv("n_updates", self._n_updates)
|
||||
|
||||
|
||||
def train_sde(self) -> None:
|
||||
# Update optimizer learning rate
|
||||
# self._update_learning_rate(self.policy.optimizer)
|
||||
|
|
|
|||
|
|
@ -1 +1 @@
|
|||
0.4.0a1
|
||||
0.4.0a2
|
||||
|
|
|
|||
Loading…
Reference in a new issue