reference/generated/tensordict.nn.distributions.AddStateIndependentNormalScale

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

AddStateIndependentNormalScale

class tensordict.nn.distributions.AddStateIndependentNormalScale(scale_shape: Size | int | tuple = None, *, scale_mapping: str = 'exp', scale_lb: Number = 0.0001, device: device | None = None, make_param: bool = True, init_value: float = 0.0)

Bases: Module

A module that returns its input as a location, with a learned scale that does not depend on the input.

The module holds one tensor, state_independent_scale, of shape scale_shape. module(loc, *others) returns (loc, scale, *others), where scale is that tensor expanded to the shape of loc, mapped by scale_mapping and clamped to at least scale_lb. Unlike NormalParamExtractor, the input is not split: all of it is the location.

Parameters:

scale_shape (Size, int or tuple of int, optional) – the shape of the scale tensor. It must match the trailing dimensions of loc, or the call raises a RuntimeError. Defaults to torch.Size(()): one scale for all the elements.

Keyword Arguments:

• scale_mapping (str, optional) – the function that maps the scale tensor to positive values, by name, for example "softplus", "exp" or "biased_softplus_1.0". See mappings(). Defaults to "exp".

• scale_lb (Number, optional) – the minimum value of the scale (the standard deviation). Defaults to 1e-4.

• device (device, optional) – the device of the scale tensor. Defaults to None, the default device.

• make_param (bool, optional) – if True, the scale tensor is a trainable Parameter. If False, it is a buffer. Defaults to True.

• init_value (float, optional) – the initial value of the scale tensor, before scale_mapping. With the default "exp" mapping, 0.0 gives a scale of 1. Defaults to 0.0.

Examples

>>> from torch import nn
>>> import torch
>>> from tensordict.nn.distributions import AddStateIndependentNormalScale
>>> num_outputs = 4
>>> module = nn.Linear(3, num_outputs)
>>> module_normal = AddStateIndependentNormalScale(num_outputs)
>>> tensor = torch.randn(3)
>>> loc, scale = module_normal(module(tensor))
>>> print(loc.shape, scale.shape)
torch.Size([4]) torch.Size([4])
>>> assert (scale > 0).all()
>>> # with modules that return more than one tensor
>>> module = nn.LSTM(3, num_outputs)
>>> module_normal = AddStateIndependentNormalScale(num_outputs)
>>> tensor = torch.randn(4, 2, 3)
>>> loc, scale, others = module_normal(*module(tensor))
>>> print(loc.shape, scale.shape)
torch.Size([4, 2, 4]) torch.Size([4, 2, 4])
>>> assert (scale > 0).all()

forward(loc: Tensor, *others: Tensor) → tuple[Tensor, ...]

Returns loc, the scale expanded to the shape of loc, and others.

Parameters:

• loc (Tensor) – the location. Its trailing dimensions must match scale_shape.

• *others – other values, returned unchanged.

Returns:

The tuple (loc, scale, *others).

Raises:

RuntimeError – if the trailing dimensions of loc do not match scale_shape.