¿Cómo retropropaga un gradiente a través de muestras aleatorias?

Estoy aprendiendo sobre políticas de gradientes y estoy teniendo dificultades para entender cómo el gradiente pasa a través de una operación aleatoria. Desde aquí: “No es posible retropropagar directamente a través de muestras aleatorias. Sin embargo, existen dos métodos principales para crear funciones sustitutas que pueden retropropagarse a través”.

Tienen un ejemplo de la “función de puntuación”:

probs = policy_network(state)

Tenga en cuenta que esto es equivalente a lo que solía llamarse multinomial

m = Categorical(probs)
action = m.sample()
next_state, reward = env.step(action)
loss = -m.log_prob(action) * reward
loss.backward()

Lo cual intenté crear un ejemplo de:

import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions import Normal
import matplotlib.pyplot as plt
from tqdm import tqdm

softplus = torch.nn.Softplus()

class Model_RL(nn.Module):
def init(self):
super(Model_RL, self).init()
self.fc1 = nn.Linear(1, 20)
self.fc2 = nn.Linear(20, 30)
self.fc3 = nn.Linear(30, 2)

def forward(self, x):
    x1 = self.fc1(x)
    x = torch.relu(x1)
    x2 = self.fc2(x)
    x = torch.relu(x2)
    x3 = softplus(self.fc3(x))
    return x3, x2, x1

básico

net_RL = Model_RL()

features = torch.tensor([1.0])
x = torch.tensor([1.0])
y = torch.tensor(3.0)

baseline = 0
baseline_lr = 0.1

epochs = 3

opt_RL = optim.Adam(net_RL.parameters(), lr=1e-3)
losses = []
xs = []
for _ in tqdm(range(epochs)):
out_RL = net_RL(x)
mu, std = out_RL[0]
dist = Normal(mu, std)
print(dist)
a = dist.sample()
log_p = dist.log_prob(a)

out = features * a
reward = -torch.square((y - out))
baseline = (1-baseline_lr)*baseline + baseline_lr*reward

loss = -(reward-baseline)*log_p

opt_RL.zero_grad()
loss.backward()
opt_RL.step()
losses.append(loss.item())

Esto parece funcionar mágicamente bien, lo cual de nuevo, no entiendo cómo el gradiente pasa a través de la operación aleatoria (pero de alguna manera lo hace).

Ahora, dado que el gradiente no puede fluir a través de la operación aleatoria, intenté reemplazar mu, std = out_RL[0] con mu, std = out_RL[0].detach() y eso provocó el error:

RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn.

Si el gradiente no pasa a través de la operación aleatoria, no entiendo por qué desvincular un tensor antes de la operación importa.