reference/generated/torchrl.objectives.two_hot_decode

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

torchrl.objectives.two_hot_decode

torchrl.objectives.two_hot_decode(logits: Tensor, bins: Tensor) → Tensor

Decode logits over a raw-value support to their scalar expectation.

Parameters:

• logits (torch.Tensor) – Categorical logits whose trailing dimension matches the support size.

• bins (torch.Tensor) – One-dimensional support in raw value space.

Returns:

The softmax-weighted expectation with the trailing category dimension removed, preserving the dtype and device of logits. Mirrored bins are paired through their probability difference before the products are summed, so an antisymmetric support such as the default symexp grid decodes uniform probabilities to exactly zero under any reduction order or fused multiply-add contraction, including the kernels that torch.compile() emits.

Examples

>>> import torch
>>> from torchrl.objectives import two_hot_decode, two_hot_encode
>>> bins = torch.tensor([-1.0, 0.0, 1.0])
>>> encoded = two_hot_encode(torch.tensor([0.25]), bins)
>>> two_hot_decode((encoded + 1e-8).log(), bins)
tensor([0.2500])