lit_sb3.utils.tensors

source module lit_sb3.utils.tensors

PyTorch tensor and parameter utilities.

Functions

• get_parameters_by_name — Extract parameters from the state dict of model.

• zip_strict — zip() function but enforces that iterables are of equal length.

• polyak_update — Perform a Polyak average update on target_params using params.

• obs_as_tensor — Moves the observation to the given device.

• flat_grad — Returns the gradients of the passed sequence of parameters into a flat gradient.

source get_parameters_by_name(model: torch.nn.Module, included_names: Iterable[str]) → list[torch.Tensor]

Extract parameters from the state dict of model.

if the name contains one of the strings in included_names.

Parameters

• model : torch.nn.Module — the model where the parameters come from.

• included_names : Iterable[str] — substrings of names to include.

Returns

• list[torch.Tensor] — List of parameters values (Pytorch tensors) that matches the queried names.

source zip_strict(*iterables: Iterable) → Iterable

zip() function but enforces that iterables are of equal length.

Raises ValueError if iterables not of equal length. Code inspired by Stackoverflow answer for question #32954486.

:param *iterables: iterables to zip()

Raises

• ValueError

source polyak_update(params: Iterable[torch.Tensor], target_params: Iterable[torch.Tensor], tau: float) → None

Perform a Polyak average update on target_params using params.

target parameters are slowly updated towards the main parameters. tau, the soft update coefficient controls the interpolation: tau=1 corresponds to copying the parameters to the target ones whereas nothing happens when tau=0. The Polyak update is done in place, with no_grad, and therefore does not create intermediate tensors, or a computation graph, reducing memory cost and improving performance. We scale the target params by 1-tau (in-place), add the new weights, scaled by tau and store the result of the sum in the target params (in place). See https://github.com/DLR-RM/stable-baselines3/issues/93.

Parameters

• params : Iterable[torch.Tensor] — parameters to use to update the target params

• target_params : Iterable[torch.Tensor] — parameters to update

• tau : float — the soft update coefficient ("Polyak update", between 0 and 1)

source obs_as_tensor(obs: np.ndarray | dict[str, np.ndarray], device: torch.device) → torch.Tensor | TensorDict

Moves the observation to the given device.

Parameters

• obs : np.ndarray | dict[str, np.ndarray] — The input observation array(s).

• device : torch.device — PyTorch device

Returns

• torch.Tensor | TensorDict — PyTorch tensor of the observation on a desired device.

Raises

• TypeError

source flat_grad(output: torch.Tensor, parameters: Sequence[nn.parameter.Parameter], create_graph: bool = False, retain_graph: bool = False) → torch.Tensor

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

Order of parameters is preserved.

Parameters

• output : torch.Tensor — functional output to compute the gradient for

• parameters : Sequence[nn.parameter.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

• torch.Tensor — Tensor containing the flattened gradients