lit_sb3.buffers.rollout
source module lit_sb3.buffers.rollout
Classes
-
RolloutBuffer — Rollout buffer used in on-policy algorithms like A2C/PPO.
source class RolloutBuffer(buffer_size: int, observation_space: spaces.Space, action_space: spaces.Space, device: torch.device | str = 'auto', gae_lambda: float = 1, gamma: float = 0.99, n_envs: int = 1)
Bases : BaseBuffer
Rollout buffer used in on-policy algorithms like A2C/PPO.
It corresponds to buffer_size transitions collected
using the current policy.
This experience will be discarded after the policy update.
In order to use PPO objective, we also store the current value of each state
and the log probability of each taken action.
The term rollout here refers to the model-free notion and should not be used with the concept of rollout used in model-based RL or planning. Hence, it is only involved in policy and value function training but not action selection.
Parameters
-
buffer_size : int — Max number of element in the buffer
-
observation_space : spaces.Space — Observation space
-
action_space : spaces.Space — Action space
-
device : torch.device | str — PyTorch device
-
gae_lambda : float — Factor for trade-off of bias vs variance for Generalized Advantage Estimator Equivalent to classic advantage when set to 1.
-
gamma : float — Discount factor
-
n_envs : int — Number of parallel environments
Methods
-
compute_returns_and_advantage — Post-processing step: compute the lambda-return (TD(lambda) estimate).
-
add — Args.
source method RolloutBuffer.reset() → None
source method RolloutBuffer.compute_returns_and_advantage(last_values: torch.Tensor, dones: np.ndarray) → None
Post-processing step: compute the lambda-return (TD(lambda) estimate).
and GAE(lambda) advantage.
Uses Generalized Advantage Estimation (https://arxiv.org/abs/1506.02438)
to compute the advantage. To obtain Monte-Carlo advantage estimate (A(s) = R - V(S))
where R is the sum of discounted reward with value bootstrap
(because we don't always have full episode), set gae_lambda=1.0 during initialization.
The TD(lambda) estimator has also two special cases
- TD(1) is Monte-Carlo estimate (sum of discounted rewards)
- TD(0) is one-step estimate with bootstrapping (r_t + gamma * v(s_{t+1}))
For more information, see discussion in https://github.com/DLR-RM/stable-baselines3/pull/375.
Parameters
-
last_values : torch.Tensor — state value estimation for the last step (one for each env)
-
dones : np.ndarray — if the last step was a terminal step (one bool for each env).
source method RolloutBuffer.add(obs: np.ndarray, action: np.ndarray, reward: np.ndarray, episode_start: np.ndarray, value: torch.Tensor, log_prob: torch.Tensor) → None
Args.
obs: Observation action: Action reward: episode_start: Start of episode signal. value: estimated value of the current state following the current policy. log_prob: log probability of the action following the current policy.
source method RolloutBuffer.get(batch_size: int | None = None) → Generator[RolloutBufferSamples, None, None]