hl-gauss-pytorch 0.1.2

HL Gauss - Pytorch

HL Gauss - Pytorch

The Gaussian Histogram Loss (HL-Gauss) proposed by Imani et al. with a few convenient wrappers, in Pytorch.

A team at Deepmind wrote a paper with a lot of positive findings for its use in value-based RL.

Install

$ pip install hl-gauss-pytorch

Usage

The HLGaussLoss module as defined in Appendix A. of the Stop Regressing paper

import torch
from hl_gauss_pytorch import HLGaussLoss

hl_gauss = HLGaussLoss(
    min_value = 0.,
    max_value = 5.,
    num_bins = 32,
    sigma = 0.5
)

logits = torch.randn(3, 16, 32).requires_grad_()
targets = torch.randint(0, 5, (3, 16)).float()

loss = hl_gauss(logits, targets)
loss.backward()

# after much training

pred_target = hl_gauss(logits) # (3, 16)

For a convenient layer that predicts from embedding to logits, import HLGaussLayer

import torch
from hl_gauss_pytorch import HLGaussLayer

hl_gauss_layer = HLGaussLayer(
    dim = 256, # input embedding dimension
    hl_gauss_loss = dict(
        min_value = 0.,
        max_value = 5.,
        num_bins = 32,
        sigma = 0.5,
    )
)

embed = torch.randn(7, 256)
targets = torch.randint(0, 5, (7,)).float()

loss = hl_gauss_layer(embed, targets)
loss.backward()

# after much training

pred_target = hl_gauss_layer(embed) # (7,)

For ablating the proposal, you can make the HLGaussLayer fall back to regular regression by setting use_regression = True, keeping the code above unchanged

HLGaussLayer(..., use_regression = True)

Todo

• use welford to keep track of running mean and std dev, in order to see if possible to dynamically set the hl gauss hparams online based only on bins + desired bin size / sigma ratio

Citations

@article{Imani2024InvestigatingTH,
    title   = {Investigating the Histogram Loss in Regression},
    author  = {Ehsan Imani and Kai Luedemann and Sam Scholnick-Hughes and Esraa Elelimy and Martha White},
    journal = {ArXiv},
    year    = {2024},
    volume  = {abs/2402.13425},
    url     = {https://api.semanticscholar.org/CorpusID:267770096}
}

@inproceedings{Imani2018ImprovingRP,
    title   = {Improving Regression Performance with Distributional Losses},
    author  = {Ehsan Imani and Martha White},
    booktitle = {International Conference on Machine Learning},
    year    = {2018},
    url     = {https://api.semanticscholar.org/CorpusID:48365278}
}

@article{Farebrother2024StopRT,
    title   = {Stop Regressing: Training Value Functions via Classification for Scalable Deep RL},
    author  = {Jesse Farebrother and Jordi Orbay and Quan Ho Vuong and Adrien Ali Taiga and Yevgen Chebotar and Ted Xiao and Alex Irpan and Sergey Levine and Pablo Samuel Castro and Aleksandra Faust and Aviral Kumar and Rishabh Agarwal},
    journal = {ArXiv},
    year   = {2024},
    volume = {abs/2403.03950},
    url    = {https://api.semanticscholar.org/CorpusID:268253088}
}