Add IQN implementation

[kashif]

Issue #, if available:

Description of changes:

Implemented IQN distribution output head from the paper https://arxiv.org/abs/2107.03743

fixes #1643

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

Please tag this pr with at least one of these labels to make our release process faster: BREAKING, new feature, bug fix, other change, dev setup

[lostella]

@kashif looks good to me, I would ask just for a few minor changes (see also the comment about pinball loss)

[lostella]

Sorry, got confused: it should be

return (self.taus - (value < self.outputs).float()) * (value - self.outputs)

Proof:

import torch
import numpy as np

def pinball_loss(obs, pred, alpha):
    return (alpha - (obs < pred).float()) * (obs - pred)

def subgradient_method(f, p, maxit):
    for k in range(maxit):
        v = f(p)
        v.backward()
        with torch.no_grad():
            p -= 1/(k+1) * p.grad
            p.grad.zero_()
    return p

data = np.random.normal(size=(1000,))
f = lambda q: pinball_loss(torch.from_numpy(data), q, 0.1).mean()
q = subgradient_method(f, torch.tensor(0.0, requires_grad=True), 1000)
print(q)

[kashif]

awesome thanks @lostella fixed

[lostella]

🚢 thanks @kashif!

[kashif]

I thank you!