Skip to content

lit_sb3.policies.base

source module lit_sb3.policies.base

Policies: abstract base class and concrete implementations.

Classes

  • BaseModel — The base model object: makes predictions in response to observations.

  • BasePolicy — The base policy object.

source class BaseModel(observation_space: spaces.Space, action_space: spaces.Space, features_extractor_class: type[BaseFeaturesExtractor] = FlattenExtractor, features_extractor_kwargs: dict[str, Any] | None = None, features_extractor: BaseFeaturesExtractor | None = None, normalize_images: bool = True, optimizer_class: type[torch.optim.Optimizer] = torch.optim.Adam, optimizer_kwargs: dict[str, Any] | None = None)

Bases : nn.Module

The base model object: makes predictions in response to observations.

In the case of policies, the prediction is an action. In the case of critics, it is the estimated value of the observation.

Parameters

  • observation_space : spaces.Space — The observation space of the environment

  • action_space : spaces.Space — The action space of the environment

  • features_extractor_class : type[BaseFeaturesExtractor] — Features extractor to use.

  • features_extractor_kwargs : dict[str, Any] | None — Keyword arguments to pass to the features extractor.

  • features_extractor : BaseFeaturesExtractor | None — Network to extract features (a CNN when using images, a nn.Flatten() layer otherwise)

  • normalize_images : bool — Whether to normalize images or not, dividing by 255.0 (True by default)

  • optimizer_class : type[torch.optim.Optimizer] — The optimizer to use, torch.optim.Adam by default

  • optimizer_kwargs : dict[str, Any] | None — Additional keyword arguments, excluding the learning rate, to pass to the optimizer

Attributes

  • device : torch.device — Infer which device this policy lives on by inspecting its parameters.

Methods

source method BaseModel.make_features_extractor() → BaseFeaturesExtractor

Helper method to create a features extractor.

source method BaseModel.extract_features(obs: PyTorchObs, features_extractor: BaseFeaturesExtractor) → torch.Tensor

Preprocess the observation if needed and extract features.

Parameters

  • obs : PyTorchObs — Observation

  • features_extractor : BaseFeaturesExtractor — The features extractor to use.

Returns

  • torch.Tensor — The extracted features

source property BaseModel.device: torch.device

Infer which device this policy lives on by inspecting its parameters.

If it has no parameters, the 'cpu' device is used as a fallback.

source method BaseModel.save(path: str) → None

Save model to a given location.

Parameters

  • path : str

  • path : str — Path to save the policy to.

source classmethod BaseModel.load(path: str, device: torch.device | str = 'auto') → SelfBaseModel

Load model from path.

Parameters

  • path : str

  • device : torch.device | str — Device on which the policy should be loaded.

  • path : str — Path to save the policy to.

source method BaseModel.load_from_vector(vector: np.ndarray) → None

Load parameters from a 1D vector.

Parameters

  • vector : np.ndarray

  • vector : np.ndarray — The vector parameter.

source method BaseModel.parameters_to_vector() → np.ndarray

Convert the parameters to a 1D vector.

source method BaseModel.set_training_mode(mode: bool) → None

Put the policy in either training or evaluation mode.

This affects certain modules, such as batch normalisation and dropout.

Parameters

  • mode : bool — if true, set to training mode, else set to evaluation mode

source method BaseModel.is_vectorized_observation(observation: np.ndarray | dict[str, np.ndarray]) → bool

Check whether or not the observation is vectorized,.

apply transposition to image (so that they are channel-first) if needed. This is used in DQN when sampling random action (epsilon-greedy policy).

Parameters

  • observation : np.ndarray | dict[str, np.ndarray] — the input observation to check

Returns

  • bool — whether the given observation is vectorized or not

source method BaseModel.obs_to_tensor(observation: np.ndarray | dict[str, np.ndarray]) → tuple[PyTorchObs, bool]

Convert an input observation to a PyTorch tensor that can be fed to a model.

Includes sugar-coating to handle different observations (e.g. normalizing images).

Parameters

  • observation : np.ndarray | dict[str, np.ndarray] — the input observation

Returns

  • tuple[PyTorchObs, bool] — The observation as PyTorch tensor and whether the observation is vectorized or not

source class BasePolicy(*args, squash_output: bool = False, **kwargs)

Bases : BaseModel, ABC

The base policy object.

Parameters are mostly the same as BaseModel; additions are documented below.

Parameters

  • args — positional arguments passed through to BaseModel.

  • kwargs — keyword arguments passed through to BaseModel.

  • squash_output : bool — For continuous actions, whether the output is squashed or not using a tanh() function.

Attributes

  • device : torch.device — Infer which device this policy lives on by inspecting its parameters.

  • squash_output : bool — (bool) Getter for squash_output.

Methods

  • init_weights — Orthogonal initialization (used in PPO and A2C).

  • predict — Get the policy action from an observation (and optional hidden state).

  • scale_action — Rescale the action from [low, high] to [-1, 1].

  • unscale_action — Rescale the action from [-1, 1] to [low, high].

source property BasePolicy.squash_output: bool

(bool) Getter for squash_output.

source staticmethod BasePolicy.init_weights(module: nn.Module, gain: float = 1) → None

Orthogonal initialization (used in PPO and A2C).

source method BasePolicy.predict(observation: np.ndarray | dict[str, np.ndarray], state: tuple[np.ndarray, ...] | None = None, episode_start: np.ndarray | None = None, deterministic: bool = False) → tuple[np.ndarray, tuple[np.ndarray, ...] | None]

Get the policy action from an observation (and optional hidden state).

Includes sugar-coating to handle different observations (e.g. normalizing images).

Parameters

  • observation : np.ndarray | dict[str, np.ndarray] — the input observation

  • state : tuple[np.ndarray, ...] | None — The last hidden states (can be None, used in recurrent policies)

  • episode_start : np.ndarray | None — The last masks (can be None, used in recurrent policies) this correspond to beginning of episodes, where the hidden states of the RNN must be reset.

  • deterministic : bool — Whether or not to return deterministic actions.

Returns

  • tuple[np.ndarray, tuple[np.ndarray, ...] | None] — the model's action and the next hidden state (used in recurrent policies)

Raises

  • ValueError

source method BasePolicy.scale_action(action: np.ndarray) → np.ndarray

Rescale the action from [low, high] to [-1, 1].

(no need for symmetric action space).

Parameters

  • action : np.ndarray — Action to scale

Returns

  • np.ndarray — Scaled action

source method BasePolicy.unscale_action(scaled_action: np.ndarray) → np.ndarray

Rescale the action from [-1, 1] to [low, high].

(no need for symmetric action space).

Parameters

  • scaled_action : np.ndarray — Action to un-scale