lit_sb3.utils.math

source module lit_sb3.utils.math

Mathematical utilities: statistics, losses, optimisation helpers.

Functions

• explained_variance — Computes fraction of variance that ypred explains about y.

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

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

• conjugate_gradient_solver — Finds an approximate solution to a set of linear equations Ax = 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 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 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 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