lit_sb3.algorithms.qrdqn

Classes

CnnPolicy — Policy class for QR-DQN when using images as input.
MultiInputPolicy — Policy class for QR-DQN when using dict observations as input.
QRDQN — Quantile Regression Deep Q-Network (QR-DQN).
QRDQNPolicy — Policy class with quantile and target networks for QR-DQN.
QuantileNetwork — Quantile network for QR-DQN.

source class CnnPolicy(observation_space: spaces.Space, action_space: spaces.Discrete, lr_schedule: Schedule, n_quantiles: int = 200, net_arch: list[int] | None = None, activation_fn: type[nn.Module] = nn.ReLU, features_extractor_class: type[BaseFeaturesExtractor] = NatureCNN, features_extractor_kwargs: dict[str, Any] | None = None, normalize_images: bool = True, optimizer_class: type[torch.optim.Optimizer] = torch.optim.Adam, optimizer_kwargs: dict[str, Any] | None = None)

Bases : QRDQNPolicy

Policy class for QR-DQN when using images as input.

Parameters

observation_space : spaces.Space — Observation space
action_space : spaces.Discrete — Action space
lr_schedule : Schedule — Learning rate schedule (could be constant)
n_quantiles : int — Number of quantiles
net_arch : list[int] | None — The specification of the network architecture.
activation_fn : type[nn.Module] — Activation function
features_extractor_class : type[BaseFeaturesExtractor] — Features extractor to use.
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.
squash_output : bool — (bool) Getter for squash_output.

source class MultiInputPolicy(observation_space: spaces.Space, action_space: spaces.Discrete, lr_schedule: Schedule, n_quantiles: int = 200, net_arch: list[int] | None = None, activation_fn: type[nn.Module] = nn.ReLU, features_extractor_class: type[BaseFeaturesExtractor] = CombinedExtractor, features_extractor_kwargs: dict[str, Any] | None = None, normalize_images: bool = True, optimizer_class: type[torch.optim.Optimizer] = torch.optim.Adam, optimizer_kwargs: dict[str, Any] | None = None)

Bases : QRDQNPolicy

Policy class for QR-DQN when using dict observations as input.

Parameters

observation_space : spaces.Space — Observation space
action_space : spaces.Discrete — Action space
lr_schedule : Schedule — Learning rate schedule (could be constant)
n_quantiles : int — Number of quantiles
net_arch : list[int] | None — The specification of the network architecture.
activation_fn : type[nn.Module] — Activation function
features_extractor_class : type[BaseFeaturesExtractor] — Features extractor to use.
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.
squash_output : bool — (bool) Getter for squash_output.

source class QRDQN(policy: str | type[QRDQNPolicy], learning_rate: float | Schedule = 5e-05, buffer_size: int = 1000000, learning_starts: int = 100, batch_size: int = 32, tau: float = 1.0, gamma: float = 0.99, train_freq: int | tuple[int, str] = 4, gradient_steps: int = 1, replay_buffer_class: type[ReplayBuffer] | None = None, replay_buffer_kwargs: dict[str, Any] | None = None, optimize_memory_usage: bool = False, n_steps: int = 1, target_update_interval: int = 10000, exploration_fraction: float = 0.005, exploration_initial_eps: float = 1.0, exploration_final_eps: float = 0.01, max_grad_norm: float | None = None, stats_window_size: int = 100, policy_kwargs: dict[str, Any] | None = None, verbose: int = 0, seed: int | None = None, device: torch.device | str = 'auto', _init_setup_model: bool = True)

Bases : OffPolicyAlgorithm

Quantile Regression Deep Q-Network (QR-DQN).

Paper: https://arxiv.org/abs/1710.10044.

Adapted from sb3-contrib to use PyTorch Lightning.

Attributes

device : torch.device — Return current device (Lightning-aware).
hparams : Union[AttributeDict, MutableMapping] — meth:save_hyperparameters. It is mutable by the user. For the frozen set of initial hyperparameters, use :attr:hparams_initial.
hparams_initial : AttributeDict — meth:save_hyperparameters. These contents are read-only. Manual updates to the saved hyperparameters can instead be performed through :attr:hparams.
example_input_array : Optional[Union[Tensor, tuple, dict]] — meth:forward method. The return type is interpreted as follows:
current_epoch : int — The current epoch in the Trainer, or 0 if not attached.
global_step : int — Total training batches seen across all epochs.
global_rank : int — The index of the current process across all nodes and devices.
local_rank : int — The index of the current process within a single node.
on_gpu : bool — Returns True if this model is currently located on a GPU.
strict_loading : bool — Determines how Lightning loads this model using .load_state_dict(..., strict=model.strict_loading).
logger : Optional[Union[Logger, FabricLogger]] — Reference to the logger object in the Trainer.
loggers : Union[list[Logger], list[FabricLogger]] — Reference to the list of loggers in the Trainer.
device_mesh : Optional['DeviceMesh'] — Strategies like ModelParallelStrategy will create a device mesh that can be accessed in the :meth:~lightning.pytorch.core.hooks.ModelHooks.configure_model hook to parallelize the LightningModule.
env : VecEnv | None — Return the active training env from the trainer, or None.

Methods

configure_optimizers
training_step — Off-policy QRDQN step: collect env transitions then do gradient updates.
predict — Epsilon-greedy action selection for QRDQN.
learn

source method QRDQN.configure_optimizers() → OptimizerLRScheduler

source method QRDQN.training_step(batch: Any, batch_idx: int) → None

Off-policy QRDQN step: collect env transitions then do gradient updates.

source method QRDQN.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]

Epsilon-greedy action selection for QRDQN.

source method QRDQN.learn(total_timesteps: int, env: GymEnv | str | None = None, callback: MaybeCallback = None, log_every_n_steps: int | None = 4, tb_log_name: str = 'QRDQN', reset_num_timesteps: bool = True, progress_bar: bool = False) → SelfQRDQN

source class QRDQNPolicy(observation_space: spaces.Space, action_space: spaces.Discrete, lr_schedule: Schedule, n_quantiles: int = 200, net_arch: list[int] | None = None, activation_fn: type[nn.Module] = nn.ReLU, features_extractor_class: type[BaseFeaturesExtractor] = FlattenExtractor, features_extractor_kwargs: dict[str, Any] | None = None, normalize_images: bool = True, optimizer_class: type[torch.optim.Optimizer] = torch.optim.Adam, optimizer_kwargs: dict[str, Any] | None = None)

Bases : BasePolicy

Policy class with quantile and target networks for QR-DQN.

Parameters

observation_space : spaces.Space — Observation space
action_space : spaces.Discrete — Action space
lr_schedule : Schedule — Learning rate schedule (could be constant)
n_quantiles : int — Number of quantiles
net_arch : list[int] | None — The specification of the network architecture.
activation_fn : type[nn.Module] — Activation function
features_extractor_class : type[BaseFeaturesExtractor] — Features extractor to use.
features_extractor_kwargs : dict[str, Any] | None — Keyword arguments to pass to the features extractor.
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.
squash_output : bool — (bool) Getter for squash_output.

Methods

make_quantile_net
forward
set_training_mode — Put the policy in either training or evaluation mode.

source method QRDQNPolicy.make_quantile_net() → QuantileNetwork

source method QRDQNPolicy.forward(obs: PyTorchObs, deterministic: bool = True) → torch.Tensor

source method QRDQNPolicy.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 class QuantileNetwork(observation_space: spaces.Space, action_space: spaces.Discrete, features_extractor: BaseFeaturesExtractor, features_dim: int, n_quantiles: int = 200, net_arch: list[int] | None = None, activation_fn: type[nn.Module] = nn.ReLU, normalize_images: bool = True)

Bases : BasePolicy

Quantile network for QR-DQN.

Parameters

observation_space : spaces.Space — Observation space
action_space : spaces.Discrete — Action space
n_quantiles : int — Number of quantiles
net_arch : list[int] | None — The specification of the network architecture.
activation_fn : type[nn.Module] — Activation function
normalize_images : bool — Whether to normalize images or not, dividing by 255.0 (True by default)

Attributes

device : torch.device — Infer which device this policy lives on by inspecting its parameters.
squash_output : bool — (bool) Getter for squash_output.

Methods

forward — Predict the quantiles.

source method QuantileNetwork.forward(obs: PyTorchObs) → torch.Tensor

Predict the quantiles.

Parameters

obs : PyTorchObs — Observation

Returns

torch.Tensor — The estimated quantiles for each action.