TanhNormal#

Implements a TanhNormal distribution with location scaling.

Location scaling prevents the location to be “too far” from 0 when a TanhTransform is applied, but ultimately leads to numerically unstable samples and poor gradient computation (e.g. gradient explosion). In practice, with location scaling the location is computed according to

\[loc = tanh(loc / upscale) * upscale.\]

Parameters:

loc (torch.Tensor) – normal distribution location parameter
scale (torch.Tensor, float, or callable) – normal distribution sigma parameter (squared root of variance). Can be a tensor, a float, or a callable that takes the loc tensor as input and returns the scale tensor. Using a callable (e.g., torch.ones_like or functools.partial(torch.full_like, fill_value=0.1)) avoids explicit device transfers like torch.tensor(val, device=device) and prevents graph breaks in torch.compile().
upscale (torch.Tensor or number) –
‘a’ scaling factor in the formula:

\[loc = tanh(loc / upscale) * upscale.\]
low (torch.Tensor or number, optional) – minimum value of the distribution. Default is -1.0;
high (torch.Tensor or number, optional) – maximum value of the distribution. Default is 1.0;
event_dims (int, optional) – number of dimensions describing the action. Default is 1. Setting event_dims to 0 will result in a log-probability that has the same shape as the input, 1 will reduce (sum over) the last dimension, 2 the last two etc.
tanh_loc (bool, optional) – if True, the above formula is used for the location scaling, otherwise the raw value is kept. Default is False;
safe_tanh (bool, optional) – if True, the Tanh transform is done “safely”, to avoid numerical overflows. This will currently break with torch.compile().

Example

>>> import torch
>>> from functools import partial
>>> from torchrl.modules.distributions import TanhNormal
>>> loc = torch.zeros(3, 4)
>>> # Using a callable scale avoids device transfers and graph breaks in torch.compile
>>> dist = TanhNormal(loc, scale=torch.ones_like)
>>> # For a custom scale value, use partial to create a callable
>>> dist = TanhNormal(loc, scale=partial(torch.full_like, fill_value=0.1))
>>> sample = dist.sample()
>>> sample.shape
torch.Size([3, 4])

get_mode()[source]#: Computes an estimation of the mode using the Adam optimizer.

property mean#: Returns the mean of the distribution.

property mode#: Returns the mode of the distribution.

rsample_and_log_prob(sample_shape: Size | tuple[int, ...] = ()) → tuple[Tensor, Tensor][source]#

Sample and score the same pre-tanh value with pathwise gradients.

Calling rsample() and log_prob() separately reconstructs the pre-tanh value from the rounded action. Once tanh saturates, that inverse can return a different Normal value and produce the wrong score and gradients. This method scores the exact Normal value used to create the action.

Parameters:: sample_shape – Leading sample dimensions.
Returns:: The action and its log probability. Their shapes are sample_shape + batch_shape + event_shape and sample_shape + batch_shape, respectively.

Use log_prob() for actions not drawn by this call.

Examples

>>> loc = torch.tensor([20.0, -20.0])
>>> scale = torch.full_like(loc, 0.1)
>>> dist = TanhNormal(loc, scale, event_dims=1)
>>> action, log_prob = dist.rsample_and_log_prob()
>>> action.shape, log_prob.shape
(torch.Size([2]), torch.Size([]))

sample_and_log_prob(sample_shape: Size | tuple[int, ...] = ()) → tuple[Tensor, Tensor][source]#

Sample and score the same pre-tanh value without pathwise gradients.

The action is detached. Its log probability keeps score-function gradients with respect to the distribution parameters.

Parameters:: sample_shape – Leading sample dimensions.
Returns:: The detached action and its differentiable log probability. Their shapes are sample_shape + batch_shape + event_shape and sample_shape + batch_shape, respectively.

property support#: Returns a Constraint object representing this distribution’s support.

TanhNormal#

Docs

Tutorials

Resources