<jupyter_start><jupyter_text>Fit a neural network to the distribution<jupyter_code>import torch
import torch.nn as nn
from torch.nn import functional as F
import torch.optim as optim
class RewardModel(nn.Module):
    def __init__(self, node_size, act_fn="relu"):
        super().__init__()
        self.act_fn = getattr(torch, act_fn)
        self.fc_1 = nn.Linear(1, node_size)
        self.fc_2 = nn.Linear(node_size, node_size)
        self.fc_3 = nn.Linear(node_size, 6)

    def forward(self, reward):
        out = self.act_fn(self.fc_1(reward))
        out = self.act_fn(self.fc_2(out))
        reward = self.fc_3(out)
        return reward
inputs = torch.Tensor(rewards).unsqueeze(1)#[:20]
states = torch.Tensor(states)#[:20]
model = RewardModel(10)
optimizer = optim.Adam(model.parameters(), lr=0.01)

num_epochs = 300

for epoch in range(num_epochs):  # loop over the dataset multiple times

    # zero the parameter gradients
    optimizer.zero_grad()

    # forward + backward + optimize
    outputs = model(inputs)
    #print('outputs before exp', outputs[0:5])
    outputs = outputs.exp()
    
    # log alpha and beta parameters. 
    loss = - torch.distributions.Beta(outputs[:,0:3], outputs[:,3:] ).log_prob(states) # criterion(outputs, labels)
    loss = loss.sum(-1).mean()
    #print( 'loss for epoch', loss.item())
    loss.backward()
    optimizer.step()

    # print statistics
    if epoch %10 ==0:
        print('epoch', epoch, 'loss is:', loss.item())
        print('first series of outputs, then states', outputs[0], states[0], inputs[0])

print('Finished Training')
res = model(torch.Tensor([-8, -6, -4, -2, -0.01]).unsqueeze(1)).exp()
print(res)
torch.distributions.Beta(res[:,:3],res[:,3:] ).sample([1])<jupyter_output><empty_output><jupyter_text>multivariate with covariance matrix<jupyter_code>np.expand_dims(rewards,1).shape
states.shape
sr = np.concatenate([states,np.expand_dims(rewards,1)], axis=1)
mus = sr.mean(axis=0)
mus
sr[:,-1].argmax()
sr[100:105]
sr[-5:]
sigmas = np.cov(sr, rowvar=False)
sigmas
import torch
torch.distributions.MultivariateNormal(torch.Tensor(mus),torch.Tensor(sigmas) ).sample()<jupyter_output><empty_output><jupyter_text>Independent<jupyter_code>torch.distributions.Normal(torch.Tensor(mus),torch.Tensor(states.std(axis=0)) ).sample()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.15420962199312716, 'fraction_numerical': 0.024484536082474227, 'mean_word_length': 4.4543325526932085, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '533568', 'license': 'MIT', 'n_tokens_mistral': 846, 'n_tokens_neox': 828, 'n_words': 162, 'path': 'notebooks/.ipynb_checkpoints/fitdensitymodel-checkpoint.ipynb', 'repo_name': 'TrentBrick/Upside-Down-Free-Energy', 'size': '2328'}