def __init__( self, observation_space: gym.spaces.Space, action_space: gym.spaces.Space, lr_schedule: Callable, use_sde: bool = False, valuehead_hidden=512, features_extractor_class: Type[ BaseFeaturesExtractor] = ResFeatureExtractor, features_extractor_kwargs: Optional[Dict[str, Any]] = None, optimizer_class: Type[th.optim.Optimizer] = th.optim.Adam, optimizer_kwargs: Optional[Dict[str, Any]] = None, ): super().__init__( observation_space=observation_space, action_space=action_space, features_extractor_class=features_extractor_class, features_extractor_kwargs=features_extractor_kwargs, optimizer_class=optimizer_class, optimizer_kwargs=optimizer_kwargs, ) self.features_extractor = features_extractor_class( self.observation_space, **self.features_extractor_kwargs) self.action_dist = CategoricalDistribution(self.action_space.n) self.valuehead_hidden = valuehead_hidden self._build(lr_schedule)
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: Type[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, normalize_images: bool = True, device: Union[th.device, str] = 'auto'): super().__init__(observation_space, action_space, features_extractor=features_extractor, normalize_images=normalize_images, device=device, squash_output=True) # Save arguments to re-create object at loading 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 action_dim = get_action_dim(self.action_space) latent_pi_net = create_mlp(features_dim, -1, net_arch, activation_fn) self.latent_pi = nn.Sequential(*latent_pi_net) last_layer_dim = net_arch[-1] if len(net_arch) > 0 else features_dim self.action_dist = CategoricalDistribution(self.action_space.n) self.mu = nn.Linear(last_layer_dim, self.action_space.n)
def test_categorical(): # The entropy can be approximated by averaging the negative log likelihood # mean negative log likelihood == entropy dist = CategoricalDistribution(N_ACTIONS) set_random_seed(1) action_logits = th.rand(N_SAMPLES, N_ACTIONS) dist = dist.proba_distribution(action_logits) actions = dist.get_actions() entropy = dist.entropy() log_prob = dist.log_prob(actions) assert th.allclose(entropy.mean(), -log_prob.mean(), rtol=2e-4)
class DiscreteActor(BasePolicy): """ Actor network (policy) for SAC. :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: (Type[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 full_std: (bool) Whether to use (n_features x n_actions) parameters for the std instead of only (n_features,) when using gSDE. :param sde_net_arch: ([int]) Network architecture for extracting features when using gSDE. If None, the latent features from the policy will be used. Pass an empty list to use the states as features. :param use_expln: (bool) Use ``expln()`` function instead of ``exp()`` when using gSDE 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 clip_mean: (float) Clip the mean output when using gSDE to avoid numerical instability. :param normalize_images: (bool) Whether to normalize images or not, dividing by 255.0 (True by default) :param device: (Union[th.device, str]) Device on which the code should run. """ 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: Type[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, normalize_images: bool = True, device: Union[th.device, str] = 'auto'): super().__init__(observation_space, action_space, features_extractor=features_extractor, normalize_images=normalize_images, device=device, squash_output=True) # Save arguments to re-create object at loading 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 action_dim = get_action_dim(self.action_space) latent_pi_net = create_mlp(features_dim, -1, net_arch, activation_fn) self.latent_pi = nn.Sequential(*latent_pi_net) last_layer_dim = net_arch[-1] if len(net_arch) > 0 else features_dim self.action_dist = CategoricalDistribution(self.action_space.n) self.mu = nn.Linear(last_layer_dim, self.action_space.n) 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. Only useful when using gSDE. It corresponds to ``th.exp(log_std)`` in the normal case, but is slightly different when using ``expln`` function (cf StateDependentNoiseDistribution doc). :return: (th.Tensor) """ msg = 'get_std() is only available when using gSDE' assert isinstance(self.action_dist, StateDependentNoiseDistribution), msg return self.action_dist.get_std(self.log_std) def reset_noise(self, batch_size: int = 1) -> None: """ Sample new weights for the exploration matrix, when using gSDE. :param batch_size: (int) """ msg = 'reset_noise() is only available when using gSDE' assert isinstance(self.action_dist, StateDependentNoiseDistribution), msg 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, 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) logits = self.mu(latent_pi) return logits, {} def forward(self, obs: th.Tensor, deterministic: bool = False) -> th.Tensor: logits, kwargs = self.get_action_dist_params(obs) # Note: the action is squashed return self.action_dist.actions_from_params(logits, deterministic=deterministic, **kwargs) def action_log_prob(self, obs: th.Tensor) -> Tuple[th.Tensor, th.Tensor]: logits, kwargs = self.get_action_dist_params(obs) # return action and associated log prob return self.action_dist.log_prob_from_params(logits, **kwargs) def _predict(self, observation: th.Tensor, deterministic: bool = False) -> th.Tensor: return self.forward(observation, deterministic)
th.tensor(0.2))) if isinstance(dist, DiagGaussianDistribution): dist = dist.proba_distribution(deterministic_actions, log_std) else: dist.sample_weights(log_std, batch_size=N_SAMPLES) dist = dist.proba_distribution(deterministic_actions, log_std, state) actions = dist.get_actions() entropy = dist.entropy() log_prob = dist.log_prob(actions) assert th.allclose(entropy.mean(), -log_prob.mean(), rtol=5e-3) categorical_params = [ (CategoricalDistribution(N_ACTIONS), N_ACTIONS), (MultiCategoricalDistribution([2, 3]), sum([2, 3])), (BernoulliDistribution(N_ACTIONS), N_ACTIONS), ] @pytest.mark.parametrize("dist, CAT_ACTIONS", categorical_params) def test_categorical(dist, CAT_ACTIONS): # The entropy can be approximated by averaging the negative log likelihood # mean negative log likelihood == entropy set_random_seed(1) action_logits = th.rand(N_SAMPLES, CAT_ACTIONS) dist = dist.proba_distribution(action_logits) actions = dist.get_actions() entropy = dist.entropy() log_prob = dist.log_prob(actions)
if isinstance(dist, DiagGaussianDistribution): dist = dist.proba_distribution(deterministic_actions, log_std) else: state = th.rand(1, N_FEATURES).repeat(N_SAMPLES, 1) dist.sample_weights(log_std, batch_size=N_SAMPLES) dist = dist.proba_distribution(deterministic_actions, log_std, state) actions = dist.get_actions() entropy = dist.entropy() log_prob = dist.log_prob(actions) assert th.allclose(entropy.mean(), -log_prob.mean(), rtol=5e-3) categorical_params = [ (CategoricalDistribution(N_ACTIONS), N_ACTIONS), (MultiCategoricalDistribution([2, 3]), sum([2, 3])), (BernoulliDistribution(N_ACTIONS), N_ACTIONS), ] @pytest.mark.parametrize("dist, CAT_ACTIONS", categorical_params) def test_categorical(dist, CAT_ACTIONS): # The entropy can be approximated by averaging the negative log likelihood # mean negative log likelihood == entropy set_random_seed(1) action_logits = th.rand(N_SAMPLES, CAT_ACTIONS) dist = dist.proba_distribution(action_logits) actions = dist.get_actions() entropy = dist.entropy() log_prob = dist.log_prob(actions)
log_std_init=th.log( th.tensor(0.2))) if isinstance(dist, DiagGaussianDistribution): dist = dist.proba_distribution(deterministic_actions, log_std) else: dist.sample_weights(log_std, batch_size=N_SAMPLES) dist = dist.proba_distribution(deterministic_actions, log_std, state) actions = dist.get_actions() entropy = dist.entropy() log_prob = dist.log_prob(actions) assert th.allclose(entropy.mean(), -log_prob.mean(), rtol=5e-3) categorical_params = [(CategoricalDistribution(N_ACTIONS), N_ACTIONS), (MultiCategoricalDistribution([2, 3]), sum([2, 3])), (BernoulliDistribution(N_ACTIONS), N_ACTIONS)] @pytest.mark.parametrize("dist, CAT_ACTIONS", categorical_params) def test_categorical(dist, CAT_ACTIONS): # The entropy can be approximated by averaging the negative log likelihood # mean negative log likelihood == entropy set_random_seed(1) action_logits = th.rand(N_SAMPLES, CAT_ACTIONS) dist = dist.proba_distribution(action_logits) actions = dist.get_actions() entropy = dist.entropy() log_prob = dist.log_prob(actions) assert th.allclose(entropy.mean(), -log_prob.mean(), rtol=5e-3)