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

Rename for consistency

Browse source 
+ add _predict to actors
+ improve sac actor code
This commit is contained in:
Antonin RAFFIN 2020-03-31 17:48:23 +02:00
parent 2bbf6a9462
commit c264403816
5 changed files with 56 additions and 44 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
tests/test_distributions.py
View file

@ -39,7 +39,7 @@ def test_squashed_gaussian(model_class):
dist = SquashedDiagGaussianDistribution(N_ACTIONS)
_, log_std = dist.proba_distribution_net(N_FEATURES)
dist = dist.proba_distribution(gaussian_mean, log_std)
actions = dist.get_action()
actions = dist.get_actions()
assert th.max(th.abs(actions)) <= 1.0
def test_sde_distribution():
@ -53,7 +53,7 @@ def test_sde_distribution():
dist.sample_weights(log_std, batch_size=N_SAMPLES)
dist = dist.proba_distribution(deterministic_actions, log_std, state)
actions = dist.get_action()
actions = dist.get_actions()
assert th.allclose(actions.mean(), dist.distribution.mean.mean(), rtol=1e-3)
assert th.allclose(actions.std(), dist.distribution.scale.mean(), rtol=1e-3)
@ -78,7 +78,7 @@ def test_entropy(dist):
dist.sample_weights(log_std, batch_size=N_SAMPLES)
dist = dist.proba_distribution(deterministic_actions, log_std, state)
actions = dist.get_action()
actions = dist.get_actions()
entropy = dist.entropy()
log_prob = dist.log_prob(actions)
assert th.allclose(entropy.mean(), -log_prob.mean(), rtol=5e-3)
@ -93,7 +93,7 @@ def test_categorical():
action_logits = th.rand(N_SAMPLES, N_ACTIONS)
dist = dist.proba_distribution(action_logits)
actions = dist.get_action()
actions = dist.get_actions()
entropy = dist.entropy()
log_prob = dist.log_prob(actions)
assert th.allclose(entropy.mean(), -log_prob.mean(), rtol=1e-4)

26
torchy_baselines/common/distributions.py
View file

@ -48,7 +48,7 @@ class Distribution(object):
"""
raise NotImplementedError
def get_action(self, deterministic: bool = False) -> th.Tensor:
def get_actions(self, deterministic: bool = False) -> th.Tensor:
"""
Return an action according to the probabilty distribution.
@ -60,7 +60,7 @@ class Distribution(object):
else:
return self.sample()
def action_from_params(self, *args, **kwargs) -> th.Tensor:
def actions_from_params(self, *args, **kwargs) -> th.Tensor:
"""
Returns a sample from the probabilty distribution
given its parameters.
@ -149,12 +149,12 @@ class DiagGaussianDistribution(Distribution):
def entropy(self) -> th.Tensor:
return sum_independent_dims(self.distribution.entropy())
def action_from_params(self, mean_actions: th.Tensor,
def actions_from_params(self, mean_actions: th.Tensor,
log_std: th.Tensor,
deterministic: bool = False) -> th.Tensor:
# Update the proba distribution
self.proba_distribution(mean_actions, log_std)
return self.get_action(deterministic=deterministic)
return self.get_actions(deterministic=deterministic)
def log_prob_from_params(self, mean_actions: th.Tensor,
log_std: th.Tensor) -> Tuple[th.Tensor, th.Tensor]:
@ -166,7 +166,7 @@ class DiagGaussianDistribution(Distribution):
:param log_std: (th.Tensor)
:return: (Tuple[th.Tensor, th.Tensor])
"""
action = self.action_from_params(mean_actions, log_std)
action = self.actions_from_params(mean_actions, log_std)
log_prob = self.log_prob(action)
return action, log_prob
@ -219,7 +219,7 @@ class SquashedDiagGaussianDistribution(DiagGaussianDistribution):
def log_prob_from_params(self, mean_actions: th.Tensor,
log_std: th.Tensor) -> Tuple[th.Tensor, th.Tensor]:
action = self.action_from_params(mean_actions, log_std)
action = self.actions_from_params(mean_actions, log_std)
log_prob = self.log_prob(action, self.gaussian_action)
return action, log_prob
@ -277,14 +277,14 @@ class CategoricalDistribution(Distribution):
def entropy(self) -> th.Tensor:
return self.distribution.entropy()
def action_from_params(self, action_logits: th.Tensor,
def actions_from_params(self, action_logits: th.Tensor,
deterministic: bool = False) -> th.Tensor:
# Update the proba distribution
self.proba_distribution(action_logits)
return self.get_action(deterministic=deterministic)
return self.get_actions(deterministic=deterministic)
def log_prob_from_params(self, action_logits: th.Tensor) -> Tuple[th.Tensor, th.Tensor]:
action = self.action_from_params(action_logits)
action = self.actions_from_params(action_logits)
log_prob = self.log_prob(action)
return action, log_prob
@ -419,7 +419,7 @@ class StateDependentNoiseDistribution(Distribution):
self.distribution = Normal(mean_actions, th.sqrt(variance + self.epsilon))
return self
def get_action(self, deterministic: bool = False) -> th.Tensor:
def get_actions(self, deterministic: bool = False) -> th.Tensor:
if deterministic:
return self.mode()
else:
@ -457,18 +457,18 @@ class StateDependentNoiseDistribution(Distribution):
return None
return sum_independent_dims(self.distribution.entropy())
def action_from_params(self, mean_actions: th.Tensor,
def actions_from_params(self, mean_actions: th.Tensor,
log_std: th.Tensor,
latent_sde: th.Tensor,
deterministic: bool = False) -> th.Tensor:
# Update the proba distribution
self.proba_distribution(mean_actions, log_std, latent_sde)
return self.get_action(deterministic=deterministic)
return self.get_actions(deterministic=deterministic)
def log_prob_from_params(self, mean_actions: th.Tensor,
log_std: th.Tensor,
latent_sde: th.Tensor) -> Tuple[th.Tensor, th.Tensor]:
action = self.action_from_params(mean_actions, log_std, latent_sde)
action = self.actions_from_params(mean_actions, log_std, latent_sde)
log_prob = self.log_prob(action)
return action, log_prob

10
torchy_baselines/ppo/policies.py
View file

@ -155,11 +155,11 @@ class PPOPolicy(BasePolicy):
"""
latent_pi, latent_vf, latent_sde = self._get_latent(obs)
# Evaluate the values for the given observations
value = self.value_net(latent_vf)
values = self.value_net(latent_vf)
distribution = self._get_action_dist_from_latent(latent_pi, latent_sde=latent_sde)
action = distribution.get_action(deterministic=deterministic)
log_prob = distribution.log_prob(action)
return action, value, log_prob
actions = distribution.get_actions(deterministic=deterministic)
log_prob = distribution.log_prob(actions)
return actions, values, log_prob
def _get_latent(self, obs: th.Tensor) -> Tuple[th.Tensor, th.Tensor, th.Tensor]:
"""
@ -212,7 +212,7 @@ class PPOPolicy(BasePolicy):
"""
latent_pi, _, latent_sde = self._get_latent(observation)
distribution = self._get_action_dist_from_latent(latent_pi, latent_sde)
return distribution.get_action(deterministic=deterministic)
return distribution.get_actions(deterministic=deterministic)
def evaluate_actions(self, obs: th.Tensor,
actions: th.Tensor) -> Tuple[th.Tensor, th.Tensor, th.Tensor]:

33
torchy_baselines/sac/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, Dict
import gym
import torch as th
@ -108,34 +108,43 @@ class Actor(BasePolicy):
'reset_noise() is only available when using SDE'
self.action_dist.sample_weights(self.log_std, batch_size=batch_size)
def get_action_dist_params(self, obs: th.Tensor) -> Tuple[th.Tensor, th.Tensor, th.Tensor]:
def get_action_dist_params(self, obs: th.Tensor) -> Tuple[th.Tensor, th.Tensor, Dict[str, th.Tensor]]:
"""
Get the parameters for the action distribution.
:param obs: (th.Tensor)
:return: (Tuple[th.Tensor, th.Tensor, Dict[str, th.Tensor]])
Mean, standard deviation and optional keyword arguments.
"""
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
mean_actions = self.mu(latent_pi)
if self.use_sde:
log_std = self.log_std
else:
latent_sde = latent_pi
if self.sde_features_extractor is not None:
latent_sde = self.sde_features_extractor(features)
return mean_actions, self.log_std, dict(latent_sde=latent_sde)
# Unstructured exploration (Original implementation)
log_std = self.log_std(latent_pi)
# Original Implementation to cap the standard deviation
log_std = th.clamp(log_std, LOG_STD_MIN, LOG_STD_MAX)
return mean_actions, log_std, latent_sde
return mean_actions, log_std, {}
def forward(self, obs: th.Tensor, deterministic: bool = False) -> th.Tensor:
mean_actions, log_std, latent_sde = self.get_action_dist_params(obs)
kwargs = dict(latent_sde=latent_sde) if self.use_sde else {}
mean_actions, log_std, kwargs = self.get_action_dist_params(obs)
# Note: the action is squashed
return self.action_dist.action_from_params(mean_actions, log_std,
return self.action_dist.actions_from_params(mean_actions, log_std,
deterministic=deterministic, **kwargs)
def action_log_prob(self, obs: th.Tensor) -> Tuple[th.Tensor, th.Tensor]:
mean_actions, log_std, latent_sde = self.get_action_dist_params(obs)
kwargs = dict(latent_sde=latent_sde) if self.use_sde else {}
mean_actions, log_std, kwargs = self.get_action_dist_params(obs)
# return action and associated log prob
return self.action_dist.log_prob_from_params(mean_actions, log_std, **kwargs)
def _predict(self, observation: th.Tensor, deterministic: bool = False) -> th.Tensor:
return self.forward(observation, deterministic)
class Critic(BasePolicy):
"""

15
torchy_baselines/td3/policies.py
View file

@ -110,7 +110,7 @@ class Actor(BasePolicy):
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]:
def evaluate_actions(self, obs: th.Tensor, actions: th.Tensor) -> Tuple[th.Tensor, th.Tensor]:
"""
Evaluate actions according to the current policy,
given the observations. Only useful when using SDE.
@ -123,7 +123,7 @@ class Actor(BasePolicy):
latent_pi, latent_sde = self._get_latent(obs)
mean_actions = self.mu(latent_pi)
distribution = self.action_dist.proba_distribution(mean_actions, self.log_std, latent_sde)
log_prob = distribution.log_prob(action)
log_prob = distribution.log_prob(actions)
return log_prob, distribution.entropy()
def reset_noise(self) -> None:
@ -149,6 +149,9 @@ class Actor(BasePolicy):
features = self.extract_features(obs)
return self.mu(features)
def _predict(self, observation: th.Tensor, deterministic: bool = False) -> th.Tensor:
return self.forward(observation, deterministic=deterministic)
class Critic(BasePolicy):
"""
@ -184,14 +187,14 @@ class Critic(BasePolicy):
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]:
def forward(self, obs: th.Tensor, actions: th.Tensor) -> Tuple[th.Tensor, th.Tensor]:
features = self.extract_features(obs)
qvalue_input = th.cat([features, action], dim=1)
qvalue_input = th.cat([features, actions], 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:
def q1_forward(self, obs: th.Tensor, actions: th.Tensor) -> th.Tensor:
features = self.extract_features(obs)
return self.q1_net(th.cat([features, action], dim=1))
return self.q1_net(th.cat([features, actions], dim=1))
class ValueFunction(BasePolicy):