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

Add docstrings

Browse source
This commit is contained in:
Antonin Raffin 2019-11-22 17:24:47 +01:00
parent f0f2f10d1e
commit c47be0086e
3 changed files with 73 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
torchy_baselines/common/base_class.py
View file

@ -23,7 +23,7 @@ class BaseRLModel(object):
:param policy_base: (BasePolicy) the base policy used by this method
:param policy_kwargs: (dict) additional arguments to be passed to the policy on creation
:param verbose: (int) the verbosity level: 0 none, 1 training information, 2 debug
:param device: (str or th.device) Device on which the code should.
:param device: (str or th.device) Device on which the code should run.
By default, it will try to use a Cuda compatible device and fallback to cpu
if it is not possible.
:param support_multi_env: (bool) Whether the algorithm supports training

29
torchy_baselines/ppo/policies.py
View file

@ -10,6 +10,20 @@ from torchy_baselines.common.distributions import make_proba_distribution,\
class PPOPolicy(BasePolicy):
"""
Policy class (with both actor and critic) for A2C and derivates (PPO).
:param observation_space: (gym.spaces.Space) Observation space
:param action_dim: (gym.spaces.Space) Action space
:param learning_rate: (callable) Learning rate schedule (could be constant)
: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: (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
"""
def __init__(self, observation_space, action_space,
learning_rate, net_arch=None, device='cpu',
activation_fn=nn.Tanh, adam_epsilon=1e-5,
@ -100,14 +114,25 @@ class PPOPolicy(BasePolicy):
return action.detach().cpu().numpy()
def evaluate_actions(self, obs, action, deterministic=False):
"""
Evaluate actions according to the current policy,
given the observations.
:param obs: (th.Tensor)
:param action: (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.
"""
latent_pi, latent_vf = self._get_latent(obs)
_, action_distribution = self._get_action_dist_from_latent(latent_pi, deterministic=deterministic)
log_prob = action_distribution.log_prob(action)
value = self.value_net(latent_vf)
return value, log_prob, action_distribution.entropy()
def value_forward(self):
pass
def value_forward(self, obs):
_, latent_vf = self._get_latent(obs)
return self.value_net(latent_vf)
MlpPolicy = PPOPolicy

45
torchy_baselines/td3/policies.py
View file

@ -6,6 +6,20 @@ from torchy_baselines.common.policies import BasePolicy, register_policy, create
class Actor(BaseNetwork):
"""
Actor network (policy) for TD3.
:param obs_dim: (int) Dimension of the observation
:param action_dim: (int) Dimension of the action space
:param net_arch: ([int]) Network architecture
: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
:param clip_noise: (float) Clip the magnitude of the noise
:param lr_sde: (float) Learning rate for the standard deviation of the noise
:param full_std: (bool) Whether to use (n_features x n_actions) parameters
for the std instead of only (n_features,)
"""
def __init__(self, obs_dim, action_dim, net_arch, activation_fn=nn.ReLU,
use_sde=False, log_std_init=-2, clip_noise=None,
lr_sde=3e-4, full_std=False):
@ -43,6 +57,15 @@ class Actor(BaseNetwork):
return th.ones((self.latent_dim, self.action_dim)).to(self.log_std.device) * self.log_std
def evaluate_actions(self, obs, action):
"""
Evaluate actions according to the current policy,
given the observations. Only useful when using SDE.
:param obs: (th.Tensor)
:param action: (th.Tensor)
:return: (th.Tensor, th.Tensor) log likelihood of taking those actions
and entropy of the action distribution.
"""
with th.no_grad():
latent_pi = self.latent_pi(obs)
mean_actions = self.actor_net(latent_pi)
@ -77,6 +100,15 @@ class Actor(BaseNetwork):
class Critic(BaseNetwork):
"""
Critic network for TD3,
in fact it represents the action-state value function (Q-value function)
:param obs_dim: (int) Dimension of the observation
:param action_dim: (int) Dimension of the action space
:param net_arch: ([int]) Network architecture
:param activation_fn: (nn.Module) Activation function
"""
def __init__(self, obs_dim, action_dim,
net_arch, activation_fn=nn.ReLU):
super(Critic, self).__init__()
@ -98,12 +130,25 @@ class Critic(BaseNetwork):
class TD3Policy(BasePolicy):
"""
Policy class (with both actor and critic) for TD3.
:param observation_space: (gym.spaces.Space) Observation space
:param action_dim: (gym.spaces.Space) Action space
:param learning_rate: (callable) Learning rate schedule (could be constant)
: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: (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
"""
def __init__(self, observation_space, action_space,
learning_rate, net_arch=None, device='cpu',
activation_fn=nn.ReLU, use_sde=False, log_std_init=-2,
clip_noise=None, lr_sde=3e-4):
super(TD3Policy, self).__init__(observation_space, action_space, device)
# Default network architecture, from the original paper
if net_arch is None:
net_arch = [400, 300]