mirror of
https://github.com/saymrwulf/stable-baselines3.git
synced 2026-07-27 20:02:30 +00:00
Remove device arg
This commit is contained in:
Antonin RAFFIN
parent
299fff67b7
commit
302e27dec4
2 changed files with 3 additions and 24 deletions
|
|
@ -1,4 +1,4 @@
|
|||
from typing import Any, Dict, List, Optional, Type, Union
|
||||
from typing import Any, Dict, List, Optional, Type
|
||||
|
||||
import gym
|
||||
import torch as th
|
||||
|
|
@ -23,7 +23,6 @@ class ActorCMA(BasePolicy):
|
|||
:param activation_fn: (Type[nn.Module]) Activation function
|
||||
: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__(
|
||||
|
|
@ -35,14 +34,12 @@ class ActorCMA(BasePolicy):
|
|||
features_dim: int,
|
||||
activation_fn: Type[nn.Module] = nn.ReLU,
|
||||
normalize_images: bool = True,
|
||||
device: Union[th.device, str] = "auto",
|
||||
):
|
||||
super(ActorCMA, self).__init__(
|
||||
observation_space,
|
||||
action_space,
|
||||
features_extractor=features_extractor,
|
||||
normalize_images=normalize_images,
|
||||
device=device,
|
||||
squash_output=True,
|
||||
)
|
||||
|
||||
|
|
@ -85,7 +82,6 @@ class CMAESPolicy(BasePolicy):
|
|||
observation_space: gym.spaces.Space,
|
||||
action_space: gym.spaces.Space,
|
||||
net_arch: Optional[List[int]] = None,
|
||||
device: Union[th.device, str] = "auto",
|
||||
activation_fn: Type[nn.Module] = nn.ReLU,
|
||||
features_extractor_class: Type[BaseFeaturesExtractor] = FlattenExtractor,
|
||||
features_extractor_kwargs: Optional[Dict[str, Any]] = None,
|
||||
|
|
@ -96,7 +92,6 @@ class CMAESPolicy(BasePolicy):
|
|||
super(CMAESPolicy, self).__init__(
|
||||
observation_space,
|
||||
action_space,
|
||||
device,
|
||||
features_extractor_class,
|
||||
features_extractor_kwargs,
|
||||
optimizer_class=optimizer_class,
|
||||
|
|
@ -124,7 +119,6 @@ class CMAESPolicy(BasePolicy):
|
|||
"net_arch": self.net_arch,
|
||||
"activation_fn": self.activation_fn,
|
||||
"normalize_images": normalize_images,
|
||||
"device": device,
|
||||
"use_sde": True,
|
||||
}
|
||||
self.actor = Actor(**self.net_args).to(self.device)
|
||||
|
|
@ -161,7 +155,6 @@ class CnnPolicy(CMAESPolicy):
|
|||
observation_space: gym.spaces.Space,
|
||||
action_space: gym.spaces.Space,
|
||||
net_arch: Optional[List[int]] = None,
|
||||
device: Union[th.device, str] = "auto",
|
||||
activation_fn: Type[nn.Module] = nn.ReLU,
|
||||
features_extractor_class: Type[BaseFeaturesExtractor] = NatureCNN,
|
||||
features_extractor_kwargs: Optional[Dict[str, Any]] = None,
|
||||
|
|
@ -173,7 +166,6 @@ class CnnPolicy(CMAESPolicy):
|
|||
observation_space,
|
||||
action_space,
|
||||
net_arch,
|
||||
device,
|
||||
activation_fn,
|
||||
features_extractor_class,
|
||||
features_extractor_kwargs,
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
|
||||
from typing import Any, Callable, Dict, List, Optional, Tuple, Type
|
||||
|
||||
import gym
|
||||
import torch as th
|
||||
|
|
@ -38,7 +38,6 @@ class Actor(BasePolicy):
|
|||
: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__(
|
||||
|
|
@ -56,14 +55,12 @@ class Actor(BasePolicy):
|
|||
use_expln: bool = False,
|
||||
clip_mean: float = 2.0,
|
||||
normalize_images: bool = True,
|
||||
device: Union[th.device, str] = "auto",
|
||||
):
|
||||
super(Actor, self).__init__(
|
||||
observation_space,
|
||||
action_space,
|
||||
features_extractor=features_extractor,
|
||||
normalize_images=normalize_images,
|
||||
device=device,
|
||||
squash_output=True,
|
||||
)
|
||||
|
||||
|
|
@ -213,7 +210,6 @@ class Critic(BaseModel):
|
|||
:param activation_fn: (Type[nn.Module]) Activation function
|
||||
: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__(
|
||||
|
|
@ -225,16 +221,14 @@ class Critic(BaseModel):
|
|||
features_dim: int,
|
||||
activation_fn: Type[nn.Module] = nn.ReLU,
|
||||
normalize_images: bool = True,
|
||||
device: Union[th.device, str] = "auto",
|
||||
n_quantiles: int = 25,
|
||||
n_critics: int = 2,
|
||||
):
|
||||
super(Critic, self).__init__(
|
||||
super().__init__(
|
||||
observation_space,
|
||||
action_space,
|
||||
features_extractor=features_extractor,
|
||||
normalize_images=normalize_images,
|
||||
device=device,
|
||||
)
|
||||
|
||||
action_dim = get_action_dim(self.action_space)
|
||||
|
|
@ -268,7 +262,6 @@ class TQCPolicy(BasePolicy):
|
|||
:param action_space: (gym.spaces.Space) Action space
|
||||
:param lr_schedule: (callable) Learning rate schedule (could be constant)
|
||||
:param net_arch: (Optional[List[int]]) 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 use_sde: (bool) Whether to use State Dependent Exploration or not
|
||||
:param log_std_init: (float) Initial value for the log standard deviation
|
||||
|
|
@ -296,7 +289,6 @@ class TQCPolicy(BasePolicy):
|
|||
action_space: gym.spaces.Space,
|
||||
lr_schedule: Callable,
|
||||
net_arch: Optional[List[int]] = None,
|
||||
device: Union[th.device, str] = "auto",
|
||||
activation_fn: Type[nn.Module] = nn.ReLU,
|
||||
use_sde: bool = False,
|
||||
log_std_init: float = -3,
|
||||
|
|
@ -314,7 +306,6 @@ class TQCPolicy(BasePolicy):
|
|||
super(TQCPolicy, self).__init__(
|
||||
observation_space,
|
||||
action_space,
|
||||
device,
|
||||
features_extractor_class,
|
||||
features_extractor_kwargs,
|
||||
optimizer_class=optimizer_class,
|
||||
|
|
@ -342,7 +333,6 @@ class TQCPolicy(BasePolicy):
|
|||
"net_arch": self.net_arch,
|
||||
"activation_fn": self.activation_fn,
|
||||
"normalize_images": normalize_images,
|
||||
"device": device,
|
||||
}
|
||||
self.actor_kwargs = self.net_args.copy()
|
||||
self.critic_kwargs = self.net_args.copy()
|
||||
|
|
@ -427,7 +417,6 @@ class CnnPolicy(TQCPolicy):
|
|||
:param action_space: (gym.spaces.Space) Action space
|
||||
:param lr_schedule: (callable) Learning rate schedule (could be constant)
|
||||
:param net_arch: (Optional[List[int]]) 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 use_sde: (bool) Whether to use State Dependent Exploration or not
|
||||
:param log_std_init: (float) Initial value for the log standard deviation
|
||||
|
|
@ -453,7 +442,6 @@ class CnnPolicy(TQCPolicy):
|
|||
action_space: gym.spaces.Space,
|
||||
lr_schedule: Callable,
|
||||
net_arch: Optional[List[int]] = None,
|
||||
device: Union[th.device, str] = "auto",
|
||||
activation_fn: Type[nn.Module] = nn.ReLU,
|
||||
use_sde: bool = False,
|
||||
log_std_init: float = -3,
|
||||
|
|
@ -473,7 +461,6 @@ class CnnPolicy(TQCPolicy):
|
|||
action_space,
|
||||
lr_schedule,
|
||||
net_arch,
|
||||
device,
|
||||
activation_fn,
|
||||
use_sde,
|
||||
log_std_init,
|
||||
|
|
|
|||
Loading…
Reference in a new issue