reference/generated/torchrl.modules.SymExpTwoHot

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

SymExpTwoHot

class torchrl.modules.SymExpTwoHot(num_bins: int = 255)

DreamerV3 categorical scalar representation.

The support contains num_bins raw scalar values obtained by applying symexp to an evenly spaced grid from -20 to 20. Targets are interpolated between adjacent raw support values, while predictions are decoded as the softmax-weighted raw-value expectation.

Parameters:

num_bins (int, optional) – Number of categorical support values. Defaults to 255.

Examples

>>> import torch
>>> from torchrl.modules import SymExpTwoHot
>>> two_hot = SymExpTwoHot(num_bins=5)
>>> target = torch.tensor([-10.0, 0.0, 10.0])
>>> encoded = two_hot.encode(target)
>>> decoded = two_hot.decode(encoded.log())
>>> torch.allclose(decoded, target, atol=1e-3)
True

decode(logits: Tensor) → Tensor

Decode categorical logits to raw scalar values.

encode(target: Tensor) → Tensor

Encode raw scalar targets as two-hot categorical targets.

forward(logits: Tensor) → Tensor

Decode logits and retain a trailing scalar event dimension.

loss(logits: Tensor, target: Tensor) → Tensor

Compute two-hot cross entropy against raw scalar targets.