lit_sb3.core.exp_manager

ExperimentManager — reads hyperparameters, creates envs, drives training.

Adapted from https://github.com/DLR-RM/rl-baselines3-zoo/blob/master/rl_zoo3/exp_manager.py

Classes

ExperimentManager — Manage a complete RL experiment: load hyperparams, create env, train and save.

source class ExperimentManager(args: Any, algo: str, env_id: str, log_folder: str, tensorboard_log: str = '', n_timesteps: int = 0, eval_freq: int = 10000, n_eval_episodes: int = 5, save_freq: int = -1, hyperparams: dict[str, Any] | None = None, env_kwargs: dict[str, Any] | None = None, eval_env_kwargs: dict[str, Any] | None = None, trained_agent: str = '', optimize_hyperparameters: bool = False, storage: str | None = None, study_name: str | None = None, n_trials: int = 1, max_total_trials: int | None = None, n_jobs: int = 1, sampler: str = 'tpe', pruner: str = 'median', optimization_log_path: str | None = None, n_startup_trials: int = 0, n_evaluations: int | None = None, truncate_last_trajectory: bool = False, uuid_str: str = '', seed: int = 0, log_interval: int = -1, save_replay_buffer: bool = False, verbose: int = 1, vec_env_type: str = 'dummy', n_eval_envs: int = 1, no_optim_plots: bool = False, device: th.device | str = 'auto', config: str | None = None, show_progress: bool = False, trial_id: int | None = None)

Manage a complete RL experiment: load hyperparams, create env, train and save.

Reads algorithm hyperparameters from a YAML file (or Python module), creates the vectorised training and eval environments, handles callbacks and optional VecNormalize, then runs model.learn() (which internally uses :class:~lit_sb3.trainer.RLTrainer).

Hyperparameter optimisation (Optuna) is supported when optuna is installed.

Parameters

args : Any — Parsed argument namespace (forwarded verbatim into saved args.yml).
algo : str — Algorithm key, e.g. "ppo".
env_id : str — Gymnasium environment ID, e.g. "CartPole-v1".
log_folder : str — Root directory for logs.
tensorboard_log : str — TensorBoard log directory ("" disabled).
n_timesteps : int — Override the number of training steps (0 use YAML).
eval_freq : int — Evaluate every n steps (<=0 no eval callback).
n_eval_episodes : int — Episodes per evaluation.
save_freq : int — Checkpoint every n steps (<=0 no checkpoints).
hyperparams : dict[str, Any] | None — Manual hyperparameter overrides.
env_kwargs : dict[str, Any] | None — Keyword arguments passed to the env constructor.
eval_env_kwargs : dict[str, Any] | None — Keyword arguments for the eval env (defaults to env_kwargs).
trained_agent : str — Path to a pre-trained .zip to continue training.
optimize_hyperparameters : bool — Run Optuna search instead of training.
storage : str | None — Optuna storage URL.
study_name : str | None — Optuna study name.
n_trials : int — Optuna trial budget per worker.
max_total_trials : int | None — Hard cap across all workers.
n_jobs : int — Parallel Optuna workers.
sampler : str — Optuna sampler ("tpe", "random", "auto").
pruner : str — Optuna pruner ("median", "halving", "none").
optimization_log_path : str | None — Where to save per-trial evaluation logs.
n_startup_trials : int — Trials before the Optuna sampler kicks in.
n_evaluations : int | None — Intermediate evaluations per trial.
truncate_last_trajectory : bool — Truncate last HER trajectory on reload.
uuid_str : str — Suffix appended to the run folder for uniqueness.
seed : int — Global random seed.
log_interval : int — Override log_every_n_steps (-1 default, <-1 disable auto-logging).
save_replay_buffer : bool — Also save the off-policy replay buffer.
verbose : int — Verbosity level (0 silent, 1 info).
vec_env_type : str — "dummy" or "subproc".
n_eval_envs : int — Parallel environments for evaluation.
no_optim_plots : bool — Skip Optuna visualisation plots.
device : th.device | str — PyTorch device.
config : str | None — Path to YAML / Python hyperparameter file (overrides default).
show_progress : bool — Display a tqdm progress bar.
trial_id : int | None — Load hyperparameters from a specific Optuna trial.

Methods

is_atari — Return True when env_id wraps an Atari game.
is_minigrid — Return True when env_id is a MiniGrid environment.
setup_experiment — Read hyperparams, create env and model.
learn — Run model.learn() with the configured parameters.
save_trained_model — Save the trained model (and optionally replay buffer + VecNormalize).
read_hyperparameters — Load and return (unprocessed_hyperparams, saved_hyperparams).
create_log_folder — Create the directory tree for this run's artefacts.
create_callbacks — Populate self.callbacks from the current configuration.
create_envs — Create, wrap and return a vectorised environment.
objective — Optuna objective function for one trial.
hyperparameters_optimization — Run the full Optuna hyperparameter search.