lit_sb3.utils

source package lit_sb3.utils

lit_sb3.utils — consolidated utility modules.

Classes

StoreDict — Custom argparse action for storing key-value pairs in a dict.
Logger — The logger class.
ConstantSchedule — Constant schedule that always returns the same value.
FloatSchedule — Wrapper that ensures the output of a Schedule is cast to float.
LinearSchedule — LinearSchedule interpolates linearly between start and end.
SimpleLinearSchedule — Linear schedule from initial_value down to 0 as training progresses.

Functions

get_device — Retrieve PyTorch device.
make_masked_proba_distribution — Return an instance of Distribution for the correct type of action space.
make_proba_distribution — Return an instance of Distribution for the correct type of action space.
create_test_env — Create a vectorised environment for evaluating a trained agent.
is_wrapped — Check if a given environment has been wrapped with a given wrapper.
make_atari_env — Create a wrapped, monitored VecEnv for Atari.
make_vec_env — Create a wrapped, monitored VecEnv.
unwrap_wrapper — Retrieve a VecEnvWrapper object by recursively searching.
configure — Configure the current logger.
get_action_masks — Checks whether gym env exposes a method returning invalid action masks.
is_masking_supported — Checks whether gym env exposes a method returning invalid action masks.
conjugate_gradient_solver — Finds an approximate solution to a set of linear equations Ax = b.
explained_variance — Computes fraction of variance that ypred explains about y.
kl_divergence — Wrapper for the PyTorch implementation of the full form KL Divergence.
quantile_huber_loss — The quantile-regression loss, as described in the QR-DQN and TQC papers.
safe_mean — Compute the mean of an array if there is at least one element.
get_hf_trained_models — Fetch trained models from a HuggingFace organisation.
get_model_path — Retrieve the path to a trained model, along with its log directory.
get_saved_hyperparams — Load hyperparameters saved by :class:~lit_sb3.core.exp_manager.ExperimentManager.
get_trained_models — Scan log_folder for trained agents saved by the ExperimentManager.
get_callback_list — Return a list of instantiated callbacks from hyperparams[\"callback\"].
get_class_by_name — Import and return a class given a dotted name string.
get_wrapper_class — Return a callable that applies one or more gym wrappers.
data_to_json — Turn data (class parameters) into a JSON string for storing.
json_to_data — Turn JSON serialization of class-parameters back into dictionary.
load_from_pkl — Load an object from the path. If a suffix is provided in the path, it will use that suffix.
load_from_zip_file — Load model data from a .zip archive.
open_path — Opens a path for reading or writing with a preferred suffix and raises debug information.
recursive_getattr — Recursive version of getattr.
recursive_setattr — Recursive version of setattr.
save_to_pkl — Save an object to path creating the necessary folders along the way.
save_to_zip_file — Save model data to a zip archive.
constant_fn — Create a function that returns a constant.
get_linear_fn — Create a function that interpolates linearly between start and end.
get_schedule_fn — Transform (if needed) learning rate and clip range (for PPO).
linear_schedule — Return a :class:SimpleLinearSchedule for initial_value.
create_sequencers — Create the utility function to chunk data into.
pad — Chunk sequences and pad them to have constant dimensions.
pad_and_flatten — Pad and flatten the sequences of scalar values,.
check_for_correct_spaces — Checks that the environment has same spaces as provided ones. Used by BaseAlgorithm to check if.
check_shape_equal — If the spaces are Box, check that they have the same shape.
is_vectorized_box_observation — For box observation type, detects and validates the shape,.
is_vectorized_dict_observation — For dict observation type, detects and validates the shape,.
is_vectorized_discrete_observation — For discrete observation type, detects and validates the shape,.
is_vectorized_multibinary_observation — For multibinary observation type, detects and validates the shape,.
is_vectorized_multidiscrete_observation — For multidiscrete observation type, detects and validates the shape,.
is_vectorized_observation — For every observation type, detects and validates the shape,.
configure_logger — Configure the logger's outputs.
get_latest_run_id — Returns the latest run number for the given log name and log path,.
get_system_info — Retrieve system and python env info for the current system.
flat_grad — Returns the gradients of the passed sequence of parameters into a flat gradient.
get_parameters_by_name — Extract parameters from the state dict of model.
obs_as_tensor — Moves the observation to the given device.
polyak_update — Perform a Polyak average update on target_params using params.
zip_strict — zip() function but enforces that iterables are of equal length.
set_random_seed — Seed the different random generators.
should_collect_more_steps — Helper used in collect_rollouts() of off-policy algorithms.
update_learning_rate — Update the learning rate for a given optimizer.

source class StoreDict(option_strings: list[str], dest: str, nargs: str = '+', **kwargs: Any)

Bases : argparse.Action

Custom argparse action for storing key-value pairs in a dict.

Enables key1:val1 key2:val2 syntax on the CLI. Numeric strings are auto-converted to int or float. Comma-separated values become lists.

source get_device(device: torch.device | str = 'auto') → torch.device

Retrieve PyTorch device.

It checks that the requested device is available first. By default, it prefers CUDA, then MPS, then CPU.

Parameters

device : torch.device | str — One of 'auto', 'cuda', 'mps', or 'cpu'

Returns

torch.device — Supported Pytorch device

source make_masked_proba_distribution(action_space: spaces.Space) → MaskableDistribution

Return an instance of Distribution for the correct type of action space.

Parameters

action_space : spaces.Space — the input action space

Returns

MaskableDistribution — the appropriate Distribution object

Raises

NotImplementedError

source make_proba_distribution(action_space: spaces.Space, use_sde: bool = False, dist_kwargs: dict[str, Any] | None = None) → Distribution

Return an instance of Distribution for the correct type of action space.

Parameters

action_space : spaces.Space — the input action space
use_sde : bool — Force the use of StateDependentNoiseDistribution instead of DiagGaussianDistribution
dist_kwargs : dict[str, Any] | None — Keyword arguments to pass to the probability distribution

Returns

Distribution — the appropriate Distribution object

Raises

NotImplementedError

Create a vectorised environment for evaluating a trained agent.

Parameters

env_id : str — Gymnasium environment ID.
n_envs : int — Number of parallel environments.
stats_path : str | Path | None — Folder containing saved VecNormalize stats.
seed : int — Random seed.
log_dir : str | Path | None — Optional directory for episode reward logging.
should_render : bool — Add render_mode="human" when True.
hyperparams : dict[str, Any] | None — Extra hyperparams (e.g. frame_stack, normalize).
env_kwargs : dict[str, Any] | None — Keyword arguments passed to the env constructor.
vec_env_cls : type[VecEnv] | None — VecEnv class to use (defaults to DummyVecEnv).
vec_env_kwargs : dict[str, Any] | None — Keyword arguments for the VecEnv constructor.

Returns

Wrapped — class:~lit_sb3.vec_env.VecEnv.

Raises

ValueError

source is_wrapped(env: gym.Env, wrapper_class: type[gym.Wrapper]) → bool

Check if a given environment has been wrapped with a given wrapper.

Parameters

env : gym.Env — Environment to check
wrapper_class : type[gym.Wrapper] — Wrapper class to look for

Returns

bool — True if environment has been wrapped with wrapper_class.

Create a wrapped, monitored VecEnv for Atari.

It is a wrapper around make_vec_env that includes common preprocessing for Atari games.

.. note:: By default, the AtariWrapper uses terminal_on_life_loss=True, which causes env.reset() to perform a no-op step instead of truly resetting when the environment terminates due to a loss of life (but not game over). To ensure reset() always resets the env, pass wrapper_kwargs=dict(terminal_on_life_loss=False).

Parameters

env_id : str | Callable[..., gym.Env] — either the env ID, the env class or a callable returning an env
n_envs : int — the number of environments you wish to have in parallel
seed : int | None — the initial seed for the random number generator
start_index : int — start rank index
monitor_dir : str | Path | None — Path to a folder where the monitor files will be saved. If None, no file will be written, however, the env will still be wrapped in a Monitor wrapper to provide additional information about training.
wrapper_kwargs : dict[str, Any] | None — Optional keyword argument to pass to the AtariWrapper
env_kwargs : dict[str, Any] | None — Optional keyword argument to pass to the env constructor
vec_env_cls : type[DummyVecEnv] | type[SubprocVecEnv] | None — A custom VecEnv class constructor. Default: None.
vec_env_kwargs : dict[str, Any] | None — Keyword arguments to pass to the VecEnv class constructor.
monitor_kwargs : dict[str, Any] | None — Keyword arguments to pass to the Monitor class constructor.

Returns

VecEnv — The wrapped environment

Create a wrapped, monitored VecEnv.

By default it uses a DummyVecEnv which is usually faster than a SubprocVecEnv.

Parameters

env_id : str | Callable[..., gym.Env] — either the env ID, the env class or a callable returning an env
n_envs : int — the number of environments you wish to have in parallel
seed : int | None — the initial seed for the random number generator
start_index : int — start rank index
monitor_dir : str | Path | None — Path to a folder where the monitor files will be saved. If None, no file will be written, however, the env will still be wrapped in a Monitor wrapper to provide additional information about training.
wrapper_class : Callable[[gym.Env], gym.Env] | None — Additional wrapper to use on the environment. This can also be a function with single argument that wraps the environment in many things. Note: the wrapper specified by this parameter will be applied after the Monitor wrapper. if some cases (e.g. with TimeLimit wrapper) this can lead to undesired behavior. See here for more details: https://github.com/DLR-RM/stable-baselines3/issues/894
env_kwargs : dict[str, Any] | None — Optional keyword argument to pass to the env constructor
vec_env_cls : type[DummyVecEnv | SubprocVecEnv] | None — A custom VecEnv class constructor. Default: None.
vec_env_kwargs : dict[str, Any] | None — Keyword arguments to pass to the VecEnv class constructor.
monitor_kwargs : dict[str, Any] | None — Keyword arguments to pass to the Monitor class constructor.
wrapper_kwargs : dict[str, Any] | None — Keyword arguments to pass to the Wrapper class constructor.

Returns

VecEnv — The wrapped environment

source unwrap_wrapper(env: gym.Env, wrapper_class: type[gym.Wrapper]) → gym.Wrapper | None

Retrieve a VecEnvWrapper object by recursively searching.

Parameters

env : gym.Env — Environment to unwrap
wrapper_class : type[gym.Wrapper] — Wrapper to look for

Returns

gym.Wrapper | None — Environment unwrapped till wrapper_class if it has been wrapped with it

source class Logger(folder: str | None, output_formats: list[KVWriter])

The logger class.

Parameters

folder : str | None — the logging location
output_formats : list[KVWriter] — the list of output formats

Methods

to_tuple — Helper function to convert str to tuple of str.
record — Log a value of some diagnostic.
record_mean — The same as record(), but if called many times, values averaged.
dump — Write all of the diagnostics from the current iteration.
log — Write the sequence of args, with no separators,.
debug — Write the sequence of args, with no separators,.
info — Write the sequence of args, with no separators,.
warn — Write the sequence of args, with no separators,.
error — Write the sequence of args, with no separators,.
set_level — Set logging threshold on current logger.
get_dir — Get directory that log files are being written to.
close — Closes the file.

source staticmethod Logger.to_tuple(string_or_tuple: str | tuple[str, ...] | None) → tuple[str, ...]

Helper function to convert str to tuple of str.

source method Logger.record(key: str, value: Any, exclude: str | tuple[str, ...] | None = None) → None

Log a value of some diagnostic.

Call this once for each diagnostic quantity, each iteration If called many times, last value will be used.

Parameters

key : str — save to log this key
value : Any — save to log this value
exclude : str | tuple[str, ...] | None — outputs to be excluded

source method Logger.record_mean(key: str, value: float | None, exclude: str | tuple[str, ...] | None = None) → None

The same as record(), but if called many times, values averaged.

Parameters

key : str — save to log this key
value : float | None — save to log this value
exclude : str | tuple[str, ...] | None — outputs to be excluded

source method Logger.dump(step: int = 0) → None

Write all of the diagnostics from the current iteration.

source method Logger.log(*args, level: int = INFO) → None

Write the sequence of args, with no separators,.

to the console and output files (if you've configured an output file).

level: int. (see logger.py docs) If the global logger level is higher than the level argument here, don't print to stdout.

Parameters

args — log the arguments
level : int — the logging level (can be DEBUG=10, INFO=20, WARN=30, ERROR=40, DISABLED=50)

source method Logger.debug(*args) → None

Write the sequence of args, with no separators,.

to the console and output files (if you've configured an output file). Using the DEBUG level.

Parameters

args — log the arguments

source method Logger.info(*args) → None

Write the sequence of args, with no separators,.

to the console and output files (if you've configured an output file). Using the INFO level.

Parameters

args — log the arguments

source method Logger.warn(*args) → None

Write the sequence of args, with no separators,.

to the console and output files (if you've configured an output file). Using the WARN level.

Parameters

args — log the arguments

source method Logger.error(*args) → None

Write the sequence of args, with no separators,.

to the console and output files (if you've configured an output file). Using the ERROR level.

Parameters

args — log the arguments

source method Logger.set_level(level: int) → None

Set logging threshold on current logger.

Parameters

level : int — the logging level (can be DEBUG=10, INFO=20, WARN=30, ERROR=40, DISABLED=50)

source method Logger.get_dir() → str | None

Get directory that log files are being written to.

will be None if there is no output directory (i.e., if you didn't call start).

Returns

str | None — the logging directory

source method Logger.close() → None

Closes the file.

source configure(folder: str | None = None, format_strings: list[str] | None = None) → Logger

Configure the current logger.

Parameters

folder : str | None — the save location (if None, $SB3_LOGDIR, if still None, tempdir/SB3-[date & time])
format_strings : list[str] | None — the output logging format (if None, $SB3_LOG_FORMAT, if still None, ['stdout', 'log', 'csv'])

Returns

Logger — The logger object.

source get_action_masks(env: GymEnv) → np.ndarray

Checks whether gym env exposes a method returning invalid action masks.

Parameters

env : GymEnv — the Gym environment to get masks from

Returns

np.ndarray — A numpy array of the masks

source is_masking_supported(env: GymEnv) → bool

Checks whether gym env exposes a method returning invalid action masks.

Parameters

env : GymEnv — the Gym environment to check

Returns

bool — True if the method is found, False otherwise

source conjugate_gradient_solver(matrix_vector_dot_fn: Callable[[torch.Tensor], torch.Tensor], b: torch.Tensor, max_iter: int = 10, residual_tol: float = 1e-10) → torch.Tensor

Finds an approximate solution to a set of linear equations Ax = b.

Sources

https://github.com/ajlangley/trpo-pytorch/blob/master/conjugate_gradient.py
https://github.com/joschu/modular_rl/blob/master/modular_rl/trpo.py#L122

Reference

https://epubs.siam.org/doi/abs/10.1137/1.9781611971446.ch6

Parameters

matrix_vector_dot_fn : Callable[[torch.Tensor], torch.Tensor] — a function that right multiplies a matrix A by a vector v
b : torch.Tensor — the right hand term in the set of linear equations Ax = b
max_iter : int — the maximum number of iterations (default is 10)
residual_tol : float — residual tolerance for early stopping of the solving (default is 1e-10) :return x: the approximate solution to the system of equations defined by matrix_vector_dot_fn and b

source explained_variance(y_pred: np.ndarray, y_true: np.ndarray) → float

Computes fraction of variance that ypred explains about y.

Returns 1 - Var[y-ypred] / Var[y].

interpretation

ev=0 => might as well have predicted zero ev=1 => perfect prediction ev<0 => worse than just predicting zero

Parameters

y_pred : np.ndarray — the prediction
y_true : np.ndarray — the expected value

Returns

float — explained variance of ypred and y

source kl_divergence(dist_true: Distribution, dist_pred: Distribution) → torch.Tensor

Wrapper for the PyTorch implementation of the full form KL Divergence.

Parameters

dist_true : Distribution — the p distribution
dist_pred : Distribution — the q distribution

Returns

torch.Tensor — KL(dist_true||dist_pred)

source quantile_huber_loss(current_quantiles: torch.Tensor, target_quantiles: torch.Tensor, cum_prob: torch.Tensor | None = None, sum_over_quantiles: bool = True) → torch.Tensor

The quantile-regression loss, as described in the QR-DQN and TQC papers.

Partially taken from https://github.com/bayesgroup/tqc_pytorch.

Parameters

current_quantiles : torch.Tensor — current estimate of quantiles, must be either (batch_size, n_quantiles) or (batch_size, n_critics, n_quantiles)
target_quantiles : torch.Tensor — target of quantiles, must be either (batch_size, n_target_quantiles), (batch_size, 1, n_target_quantiles), or (batch_size, n_critics, n_target_quantiles)
cum_prob : torch.Tensor | None — cumulative probabilities to calculate quantiles (also called midpoints in QR-DQN paper), must be either (batch_size, n_quantiles), (batch_size, 1, n_quantiles), or (batch_size, n_critics, n_quantiles). (if None, calculating unit quantiles)
sum_over_quantiles : bool — if summing over the quantile dimension or not

Returns

torch.Tensor — the loss

Raises

ValueError

source safe_mean(arr: np.ndarray | list | deque) → float

Compute the mean of an array if there is at least one element.

For empty array, return NaN. It is used for logging only.

Parameters

arr : np.ndarray | list | deque — Numpy array or list of values

source get_hf_trained_models(organization: str = 'sb3') → dict[str, tuple[str, str]]

Fetch trained models from a HuggingFace organisation.

Returns a dict mapping model names to (algo, env_id) tuples, mirroring :func:get_trained_models.

Requires huggingface_hub to be installed.

:param organization: HuggingFace organisation to search. :return: Dict mapping model names to (algo, env_id) tuples.

Raises

ImportError

source get_model_path(exp_id: int, folder: str, algo: str, env_name: EnvironmentName, load_best: bool = False, load_checkpoint: str | None = None, load_last_checkpoint: bool = False) → tuple[str, str, str]

Retrieve the path to a trained model, along with its log directory.

Parameters

exp_id : int — Experiment ID (0 → use the latest available run).
folder : str — Root folder where run directories are stored (e.g. logs/).
algo : str — Algorithm name, e.g. "ppo".
env_name : EnvironmentName — :class:EnvironmentName wrapper for the target environment.
load_best : bool — When True, load best_model.zip instead of the final checkpoint.
load_checkpoint : str | None — If non-None, load the model saved at this specific step count (e.g. "10000").
load_last_checkpoint : bool — When True, load the numerically largest step-count checkpoint rather than the final model.

Returns

tuple[str, str, str] — (name_prefix, model_path, log_path) where name_prefix is a human-readable identifier, model_path points to the .zip file, and log_path is the run directory.

Raises

ValueError

source get_saved_hyperparams(stats_path: str | Path, norm_reward: bool = False, test_mode: bool = False) → tuple[dict[str, Any], str | Path | None]

Load hyperparameters saved by :class:~lit_sb3.core.exp_manager.ExperimentManager.

Parameters

stats_path : str | Path — Path to the folder containing config.yml.
norm_reward : bool — Whether to keep reward normalisation active.
test_mode : bool — Disable reward normalisation and training flag.

Returns

tuple[dict[str, Any], str | Path | None] — Tuple of (hyperparams_dict, stats_path_or_None).

source get_trained_models(log_folder: str) → dict[str, tuple[str, str]]

Scan log_folder for trained agents saved by the ExperimentManager.

Parameters

log_folder : str — Root log directory (same as --log-folder in training).

Returns

dict[str, tuple[str, str]] — Dict mapping ModelName strings to (algo, env_id) tuples.

source get_callback_list(hyperparams: dict[str, Any]) → list[BaseCallback]

Return a list of instantiated callbacks from hyperparams[\"callback\"].

Each entry can be

a dotted class name string,
{ClassName: kwargs_dict}, or
an already-instantiated :class:~lit_sb3.callbacks.BaseCallback.

source get_class_by_name(name: str) → type

Import and return a class given a dotted name string.

Example::

get_class_by_name("lit_sb3.callbacks.CheckpointCallback")

source get_wrapper_class(hyperparams: dict[str, Any], key: str = 'env_wrapper') → Callable[[gym.Env], gym.Env] | None

Return a callable that applies one or more gym wrappers.

The key entry in hyperparams may be a single dotted class name or a list. Each entry may optionally be a {ClassName: kwargs_dict} mapping to pass keyword arguments to the wrapper constructor.

Parameters

hyperparams : dict[str, Any] — Hyperparameter dict that may contain key.
key : str — The hyperparams key to inspect ("env_wrapper" or "vec_env_wrapper").

Returns

Callable[[gym.Env], gym.Env] | None — A callable wrap_env(env) -> env or None if the key is absent.

Raises

TypeError

source data_to_json(data: dict[str, Any]) → str

Turn data (class parameters) into a JSON string for storing.

Parameters

data : dict[str, Any] — Dictionary of class parameters to be stored. Items that are not JSON serializable will be pickled with Cloudpickle and stored as bytearray in the JSON file

Returns

str — JSON string of the data serialized.

source json_to_data(json_string: str, custom_objects: dict[str, Any] | None = None) → dict[str, Any]

Turn JSON serialization of class-parameters back into dictionary.

Parameters

json_string : str — JSON serialization of the class-parameters that should be loaded.
custom_objects : dict[str, Any] | None — Dictionary of objects to replace upon loading. If a variable is present in this dictionary as a key, it will not be deserialized and the corresponding item will be used instead. Similar to custom_objects in keras.models.load_model. Useful when you have an object in file that can not be deserialized.

Returns

dict[str, Any] — Loaded class parameters.

Raises

ValueError

source load_from_pkl(path: str | pathlib.Path | io.BufferedIOBase, verbose: int = 0) → Any

Load an object from the path. If a suffix is provided in the path, it will use that suffix.

If the path does not exist, it will attempt to load using the .pkl suffix.

Parameters

path : str | pathlib.Path | io.BufferedIOBase — the path to open. if save_path is a str or pathlib.Path and mode is "w", single dispatch ensures that the path actually exists. If path is a io.BufferedIOBase the path exists.
verbose : int — Verbosity level: 0 for no output, 1 for info messages, 2 for debug messages

Load model data from a .zip archive.

Parameters

load_path : str | pathlib.Path | io.BufferedIOBase — Where to load the model from
load_data : bool — Whether we should load and return data (class parameters). Mainly used by 'load_parameters' to only load model parameters (weights)
custom_objects : dict[str, Any] | None — Dictionary of objects to replace upon loading. If a variable is present in this dictionary as a key, it will not be deserialized and the corresponding item will be used instead. Similar to custom_objects in keras.models.load_model. Useful when you have an object in file that can not be deserialized.
device : torch.device | str — Device on which the code should run.
verbose : int — Verbosity level: 0 for no output, 1 for info messages, 2 for debug messages
print_system_info : bool — Whether to print or not the system info about the saved model.

Returns

tuple[dict[str, Any] | None, TensorDict, TensorDict | None] — Class parameters, model state_dicts (aka "params", dict of state_dict) and dict of pytorch variables

Raises

ValueError

Opens a path for reading or writing with a preferred suffix and raises debug information.

If the provided path is a derivative of io.BufferedIOBase it ensures that the file matches the provided mode, i.e. If the mode is read ("r", "read") it checks that the path is readable. If the mode is write ("w", "write") it checks that the file is writable.

If the provided path is a string or a pathlib.Path, it ensures that it exists. If the mode is "read" it checks that it exists, if it doesn't exist it attempts to read path.suffix if a suffix is provided. If the mode is "write" and the path does not exist, it creates all the parent folders. If the path points to a folder, it changes the path to path_2. If the path already exists and verbose >= 2, it raises a warning.

Parameters

path : str | pathlib.Path | io.BufferedIOBase — the path to open. if save_path is a str or pathlib.Path and mode is "w", single dispatch ensures that the path actually exists. If path is a io.BufferedIOBase the path exists.
mode : str — how to open the file. "w"|"write" for writing, "r"|"read" for reading.
verbose : int — Verbosity level: 0 for no output, 1 for info messages, 2 for debug messages
suffix : str | None — The preferred suffix. If mode is "w" then the opened file has the suffix. If mode is "r" then we attempt to open the path. If an error is raised and the suffix is not None, we attempt to open the path with the suffix.

Raises

TypeError
ValueError

source recursive_getattr(obj: Any, attr: str, *args) → Any

Recursive version of getattr.

taken from https://stackoverflow.com/questions/31174295.

Ex

MyObject.sub_object = SubObject(name='test') recursive_getattr(MyObject, 'sub_object.name') # return test

Parameters

obj : Any — Object to retrieve attribute from.
attr : str — Attribute to retrieve.
*args — Default value if attribute is not found.

Returns

Any — The attribute

source recursive_setattr(obj: Any, attr: str, val: Any) → None

Recursive version of setattr.

taken from https://stackoverflow.com/questions/31174295.

Ex

MyObject.sub_object = SubObject(name='test') recursive_setattr(MyObject, 'sub_object.name', 'hello')

Parameters

obj : Any
attr : str — Attribute to set
val : Any — New value of the attribute
obj : Any — The obj parameter.

source save_to_pkl(path: str | pathlib.Path | io.BufferedIOBase, obj: Any, verbose: int = 0) → None

Save an object to path creating the necessary folders along the way.

If the path exists and is a directory, it will raise a warning and rename the path. If a suffix is provided in the path, it will use that suffix, otherwise, it will use '.pkl'.

Parameters

path : str | pathlib.Path | io.BufferedIOBase — the path to open. if save_path is a str or pathlib.Path and mode is "w", single dispatch ensures that the path actually exists. If path is a io.BufferedIOBase the path exists.
obj : Any — The object to save.
verbose : int — Verbosity level: 0 for no output, 1 for info messages, 2 for debug messages

Save model data to a zip archive.

Parameters

save_path : str | pathlib.Path | io.BufferedIOBase — Where to store the model. if save_path is a str or pathlib.Path ensures that the path actually exists.
data : dict[str, Any] | None — Class parameters being stored (non-PyTorch variables)
params : dict[str, Any] | None — Model parameters being stored expected to contain an entry for every state_dict with its name and the state_dict.
pytorch_variables : dict[str, Any] | None — Other PyTorch variables expected to contain name and value of the variable.
verbose : int — Verbosity level: 0 for no output, 1 for info messages, 2 for debug messages

source class ConstantSchedule(val: float)

Constant schedule that always returns the same value.

Useful for fixed learning rates or clip ranges.

Parameters

val : float — constant value

source class FloatSchedule(value_schedule: Schedule | float)

Wrapper that ensures the output of a Schedule is cast to float.

Can wrap either a constant value or an existing callable Schedule.

Parameters

value_schedule : Schedule | float — Constant value or callable schedule (e.g. LinearSchedule, ConstantSchedule)

source class LinearSchedule(start: float, end: float, end_fraction: float)

LinearSchedule interpolates linearly between start and end.

between progress_remaining = 1 and progress_remaining = end_fraction. This is used in DQN for linearly annealing the exploration fraction (epsilon for the epsilon-greedy strategy).

Parameters

start : float — value to start with if progress_remaining = 1
end : float — value to end with if progress_remaining = 0
end_fraction : float — fraction of progress_remaining where end is reached e.g 0.1 then end is reached after 10% of the complete training process.

source class SimpleLinearSchedule(initial_value: float | str)

Linear schedule from initial_value down to 0 as training progresses.

Parameters

initial_value : float | str — Starting value (must be positive).

source constant_fn(val: float) → Schedule

Create a function that returns a constant.

It is useful for learning rate schedule (to avoid code duplication).

Parameters

val : float — constant value

Returns

Schedule — Constant schedule function.

source get_linear_fn(start: float, end: float, end_fraction: float) → Schedule

Create a function that interpolates linearly between start and end.

between progress_remaining = 1 and progress_remaining = end_fraction. This is used in DQN for linearly annealing the exploration fraction (epsilon for the epsilon-greedy strategy).

:params start: value to start with if progress_remaining = 1 :params end: value to end with if progress_remaining = 0 :params end_fraction: fraction of progress_remaining where end is reached e.g 0.1 then end is reached after 10% of the complete training process.

Returns

Schedule — Linear schedule function.

source get_schedule_fn(value_schedule: Schedule | float) → Schedule

Transform (if needed) learning rate and clip range (for PPO).

to callable.

Parameters

value_schedule : Schedule | float — Constant value of schedule function

Returns

Schedule — Schedule function (can return constant value)

source linear_schedule(initial_value: float | str) → SimpleLinearSchedule

Return a :class:SimpleLinearSchedule for initial_value.

source create_sequencers(episode_starts: np.ndarray, env_change: np.ndarray, device: torch.device) → tuple[np.ndarray, Callable, Callable]

Create the utility function to chunk data into.

sequences and pad them to create fixed size tensors.

Parameters

episode_starts : np.ndarray — Indices where an episode starts
env_change : np.ndarray — Indices where the data collected come from a different env (when using multiple env for data collection)
device : torch.device — PyTorch device

Returns

tuple[np.ndarray, Callable, Callable] — Indices of the transitions that start a sequence, pad and pad_and_flatten utilities tailored for this batch (sequence starts and ends indices are fixed)

source pad(seq_start_indices: np.ndarray, seq_end_indices: np.ndarray, device: torch.device, tensor: np.ndarray, padding_value: float = 0.0) → torch.Tensor

Chunk sequences and pad them to have constant dimensions.

Parameters

seq_start_indices : np.ndarray — Indices of the transitions that start a sequence
seq_end_indices : np.ndarray — Indices of the transitions that end a sequence
device : torch.device — PyTorch device
tensor : np.ndarray — Tensor of shape (batch_size, *tensor_shape)
padding_value : float — Value used to pad sequence to the same length (zero padding by default)

Returns

torch.Tensor — (n_seq, max_length, *tensor_shape)

source pad_and_flatten(seq_start_indices: np.ndarray, seq_end_indices: np.ndarray, device: torch.device, tensor: np.ndarray, padding_value: float = 0.0) → torch.Tensor

Pad and flatten the sequences of scalar values,.

while keeping the sequence order. From (batch_size, 1) to (n_seq, max_length, 1) -> (n_seq * max_length,).

Parameters

seq_start_indices : np.ndarray — Indices of the transitions that start a sequence
seq_end_indices : np.ndarray — Indices of the transitions that end a sequence
device : torch.device — PyTorch device (cpu, gpu, ...)
tensor : np.ndarray — Tensor of shape (max_length, n_seq, 1)
padding_value : float — Value used to pad sequence to the same length (zero padding by default)

Returns

torch.Tensor — (n_seq * max_length,) aka (padded_batch_size,)

source check_for_correct_spaces(env: GymEnv, observation_space: spaces.Space, action_space: spaces.Space) → None

Checks that the environment has same spaces as provided ones. Used by BaseAlgorithm to check if.

spaces match after loading the model with given env. Checked parameters: - observation_space - action_space.

Parameters

env : GymEnv — Environment to check for valid spaces
observation_space : spaces.Space — Observation space to check against
action_space : spaces.Space — Action space to check against

Raises

ValueError

source check_shape_equal(space1: spaces.Space, space2: spaces.Space) → None

If the spaces are Box, check that they have the same shape.

If the spaces are Dict, it recursively checks the subspaces.

Parameters

space1 : spaces.Space — Space
space2 : spaces.Space — Other space

source is_vectorized_box_observation(observation: np.ndarray, observation_space: spaces.Box) → bool

For box observation type, detects and validates the shape,.

then returns whether or not the observation is vectorized.

Parameters

observation : np.ndarray — the input observation to validate
observation_space : spaces.Box — the observation space

Returns

bool — whether the given observation is vectorized or not

Raises

ValueError

source is_vectorized_dict_observation(observation: np.ndarray, observation_space: spaces.Dict) → bool

For dict observation type, detects and validates the shape,.

then returns whether or not the observation is vectorized.

Parameters

observation : np.ndarray — the input observation to validate
observation_space : spaces.Dict — the observation space

Returns

bool — whether the given observation is vectorized or not

Raises

ValueError

source is_vectorized_discrete_observation(observation: int | np.ndarray, observation_space: spaces.Discrete) → bool

For discrete observation type, detects and validates the shape,.

then returns whether or not the observation is vectorized.

Parameters

observation : int | np.ndarray — the input observation to validate
observation_space : spaces.Discrete — the observation space

Returns

bool — whether the given observation is vectorized or not

Raises

ValueError

source is_vectorized_multibinary_observation(observation: np.ndarray, observation_space: spaces.MultiBinary) → bool

For multibinary observation type, detects and validates the shape,.

then returns whether or not the observation is vectorized.

Parameters

observation : np.ndarray — the input observation to validate
observation_space : spaces.MultiBinary — the observation space

Returns

bool — whether the given observation is vectorized or not

Raises

ValueError

source is_vectorized_multidiscrete_observation(observation: np.ndarray, observation_space: spaces.MultiDiscrete) → bool

For multidiscrete observation type, detects and validates the shape,.

then returns whether or not the observation is vectorized.

Parameters

observation : np.ndarray — the input observation to validate
observation_space : spaces.MultiDiscrete — the observation space

Returns

bool — whether the given observation is vectorized or not

Raises

ValueError

source is_vectorized_observation(observation: int | np.ndarray, observation_space: spaces.Space) → bool

For every observation type, detects and validates the shape,.

then returns whether or not the observation is vectorized.

Parameters

observation : int | np.ndarray — the input observation to validate
observation_space : spaces.Space — the observation space

Returns

bool — whether the given observation is vectorized or not

Raises

ValueError

source configure_logger(verbose: int = 0, tensorboard_log: str | Path | None = None, tb_log_name: str = '', reset_num_timesteps: bool = True) → Logger

Configure the logger's outputs.

Parameters

verbose : int — Verbosity level: 0 for no output, 1 for the standard output to be part of the logger outputs
tensorboard_log : str | Path | None — the log location for tensorboard (if None, no logging)
tb_log_name : str — tensorboard log
reset_num_timesteps : bool — Whether the num_timesteps attribute is reset or not. It allows to continue a previous learning curve (reset_num_timesteps=False) or start from t=0 (reset_num_timesteps=True, the default).

Returns

Logger — The logger object

Raises

ImportError

source get_latest_run_id(log_path: str | Path = '', log_name: str = '') → int

Returns the latest run number for the given log name and log path,.

by finding the greatest number in the directories.

Parameters

log_path : str | Path — Path to the log folder containing several runs.
log_name : str — Name of the experiment. Each run is stored in a folder named log_name_1, log_name_2, ...

Returns

int — latest run number

source get_system_info(print_info: bool = True) → tuple[dict[str, str], str]

Retrieve system and python env info for the current system.

Parameters

print_info : bool — Whether to print or not those infos

Returns

tuple[dict[str, str], str] — Dictionary summing up the version for each relevant package and a formatted string.

source flat_grad(output: torch.Tensor, parameters: Sequence[nn.parameter.Parameter], create_graph: bool = False, retain_graph: bool = False) → torch.Tensor

Returns the gradients of the passed sequence of parameters into a flat gradient.

Order of parameters is preserved.

Parameters

output : torch.Tensor — functional output to compute the gradient for
parameters : Sequence[nn.parameter.Parameter] — sequence of Parameter
retain_graph : bool — If False, the graph used to compute the grad will be freed. Defaults to the value of create_graph.
create_graph : bool — If True, graph of the derivative will be constructed, allowing to compute higher order derivative products. Default: False.

Returns

torch.Tensor — Tensor containing the flattened gradients

source get_parameters_by_name(model: torch.nn.Module, included_names: Iterable[str]) → list[torch.Tensor]

Extract parameters from the state dict of model.

if the name contains one of the strings in included_names.

Parameters

model : torch.nn.Module — the model where the parameters come from.
included_names : Iterable[str] — substrings of names to include.

Returns

list[torch.Tensor] — List of parameters values (Pytorch tensors) that matches the queried names.

source obs_as_tensor(obs: np.ndarray | dict[str, np.ndarray], device: torch.device) → torch.Tensor | TensorDict

Moves the observation to the given device.

Parameters

obs : np.ndarray | dict[str, np.ndarray] — The input observation array(s).
device : torch.device — PyTorch device

Returns

torch.Tensor | TensorDict — PyTorch tensor of the observation on a desired device.

Raises

TypeError

source polyak_update(params: Iterable[torch.Tensor], target_params: Iterable[torch.Tensor], tau: float) → None

Perform a Polyak average update on target_params using params.

target parameters are slowly updated towards the main parameters. tau, the soft update coefficient controls the interpolation: tau=1 corresponds to copying the parameters to the target ones whereas nothing happens when tau=0. The Polyak update is done in place, with no_grad, and therefore does not create intermediate tensors, or a computation graph, reducing memory cost and improving performance. We scale the target params by 1-tau (in-place), add the new weights, scaled by tau and store the result of the sum in the target params (in place). See https://github.com/DLR-RM/stable-baselines3/issues/93.

Parameters

params : Iterable[torch.Tensor] — parameters to use to update the target params
target_params : Iterable[torch.Tensor] — parameters to update
tau : float — the soft update coefficient ("Polyak update", between 0 and 1)

source zip_strict(*iterables: Iterable) → Iterable

zip() function but enforces that iterables are of equal length.

Raises ValueError if iterables not of equal length. Code inspired by Stackoverflow answer for question #32954486.

:param *iterables: iterables to zip()

Raises

ValueError

source set_random_seed(seed: int, using_cuda: bool = False) → None

Seed the different random generators.

Parameters

seed : int — Random seed for reproducibility.
using_cuda : bool — Whether CUDA is being used.

source should_collect_more_steps(train_freq: TrainFreq, num_collected_steps: int, num_collected_episodes: int) → bool

Helper used in collect_rollouts() of off-policy algorithms.

to determine the termination condition.

Parameters

train_freq : TrainFreq — How much experience should be collected before updating the policy.
num_collected_steps : int — The number of already collected steps.
num_collected_episodes : int — The number of already collected episodes.

Returns

bool — Whether to continue or not collecting experience by doing rollouts of the current policy.

Raises

ValueError

source update_learning_rate(optimizer: torch.optim.Optimizer, learning_rate: float) → None

Update the learning rate for a given optimizer.

Useful when doing linear schedule.

Parameters

optimizer : torch.optim.Optimizer — Pytorch optimizer
learning_rate : float — New learning rate value