[Python] No convergence while trying to train sampler using RL

I am trying to replicate the technique from this paper (https://arxiv.org/pdf/2305.13301) on a simpler case. I want to fit the gaussian mean using the RL technique, but the algorithm fails to converge due to small gradients. What do i miss?

import torch
from torch.optim import SGD
from rewards import mahalanobis_reward, euclidean_centroid_distance
import numpy as np
from utils import save_gradients, set_zero_grad, add_gradients, plot_distance_to_dataset


class GaussianSampler(torch.nn.Module):
def __init__(self):
super(GaussianSampler, self).__init__()
self.mu = torch.nn.Parameter(torch.randn((2,), requires_grad=True))
self.sigma = torch.eye(2)

def forward(self, sample_size):
return torch.matmul(torch.randn((sample_size, 2)), self.sigma) + torch.stack([self.mu]*sample_size)


def rl_gaussian_fitting(epochs, model, test_dataset):
optimizer = torch.optim.AdamW(model.parameters(), lr=0.1)
sample_size = 100
for epoch in range(epochs):
optimizer.zero_grad()

samples = model(sample_size)
np_samples = samples.detach().numpy()

log_prob = torch.nn.MSELoss()
test_data = test_dataset[:sample_size, 0:2]
test_data = torch.from_numpy(test_data)
test_data = test_data.type(torch.FloatTensor)
losses = log_prob(torch.stack([model.mu]*sample_size), samples)

reward = mahalanobis_reward(np_samples, test_dataset)

#loss = log_prob * reward
loss = losses * reward
loss.backward()

optimizer.step()

if epoch % 10 == 0:
print('Epoch {}/{}, Loss: {:.4f}'.format(epoch, epochs, loss.mean().item()))
print('Reward: {:.4f}'.format(reward))
print('mu: {}'.format(model.mu.data))
print('mu_grad: {}'.format(model.mu.grad.data))

plot_distance_to_dataset(np_samples, test_dataset, 'Distance_Gauss'+'_reward_'+str(reward))


if __name__ == '__main__':
model = GaussianSampler()
test_dataset = np.load('data/one_cloud_dataset.npy')
rl_gaussian_fitting(1000, model, test_dataset)
print(model.mu, model.sigma)


the test dataset is just 100 samples from N((-10,10), I)

Continue reading...