Utils

sb3_contrib.common.utils.conjugate_gradient_solver(matrix_vector_dot_fn, b, max_iter=10, residual_tol=1e-10)

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[[Tensor], Tensor]) – a function that right multiplies a matrix A by a vector v

• b – the right hand term in the set of linear equations Ax = b

• max_iter – the maximum number of iterations (default is 10)

• residual_tol – 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

Return type:

Tensor

sb3_contrib.common.utils.flat_grad(output, parameters, create_graph=False, retain_graph=False)

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

Parameters:

• output – functional output to compute the gradient for

• parameters (Sequence[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:

Tensor containing the flattened gradients

Return type:

Tensor

sb3_contrib.common.utils.quantile_huber_loss(current_quantiles, target_quantiles, cum_prob=None, sum_over_quantiles=True)

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 (Tensor) – current estimate of quantiles, must be either (batch_size, n_quantiles) or (batch_size, n_critics, n_quantiles)

• target_quantiles (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 (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:

the loss

Return type:

Tensor