lit_sb3.utils.training

source module lit_sb3.utils.training

Training utilities: seeding, learning rate, rollout control.

Functions

• set_random_seed — Seed the different random generators.

• update_learning_rate — Update the learning rate for a given optimizer.

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

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 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

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