<jupyter_start><jupyter_text>$$\mathbb{E}_{p(b|\theta)} \left[ (b - .45)^2 \right]$$<jupyter_code>def logprob(b, theta):
    #return (b - theta)**2
    
    probs = torch.nn.functional.log_softmax(theta)
    b_ohe = torch.zeros_like(probs)
    b_ohe[b] = 1
    return torch.dot(b_ohe, probs)
    #return b*torch.log(theta) + (1 - b)*torch.log(1 - theta)
    
torch.multinomial(torch.tensor([0.5, 0.5]), num_samples=1)[0]
torch.bernoulli(torch.tensor(0.5))
def magic_box(x):
    return torch.exp(x - x.detach())


def objective_reinforce(fun, logprob):
    def f(theta, b):
        #b = torch.bernoulli(theta)
        #bb = magic_box(b)
        return fun(b) * magic_box(logprob(b, theta))
    return f


def fun(b):
    return (b - .1)**2
theta0 = torch.tensor([0.5, 0.5], requires_grad=True)
eta = torch.tensor(0.005, requires_grad=True)

def reinforce(theta0, eta):
    """One step of reinforce gradient estimator
    """
    theta_p = theta0
    for i in range(1000):
        #bs = 0.1 * torch.randn(1) + theta_p
        #bs = torch.bernoulli(theta_p)
        bs = torch.multinomial(theta_p, num_samples=1)[0]
        theta_p = theta_p - eta*grad(objective_reinforce(fun, logprob)(theta_p, bs), theta_p, create_graph=True)[0]
    #loss = fun( 0.1 * torch.randn(1) + theta_p).mean()
    loss = fun( torch.multinomial(theta_p, num_samples=1) ).mean()
    #loss = fun( torch.bernoulli(theta_p) ).mean()
    #losses.append(loss)   
    return loss, theta_p
reinforce(theta0, eta)
grad(reinforce(theta0, eta)[0], theta0)[0]
theta0 = torch.tensor(.25, requires_grad=True)
eta = torch.tensor(0.1, requires_grad=True)


thetas = []

for _ in range(1000):

    eta = eta - 0.01*grad(reinforce(theta0, eta), eta)[0]
    thetas.append(eta)
    #print(grad(reinforce(theta0, eta), theta0)[0])



plt.plot(thetas)<jupyter_output><empty_output><jupyter_text>Simple higher order autodiff example<jupyter_code>def loss(x):
    return torch.sum(x**2)

def sgd(x0, eta):
    return loss(x0 - eta*grad(loss(x0), x0)[0])
loss(torch.tensor(0.2))
sgd(torch.tensor(0.2, requires_grad=True), eta=0.01)
x0 = torch.tensor(0.2, requires_grad=True)
eta = torch.tensor(0.01, requires_grad=True)

grad(sgd(x0, eta), eta)
grad(sgd(x0, eta), x0)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1705253784505788, 'fraction_numerical': 0.03829029385574354, 'mean_word_length': 4.00445434298441, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 9, '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': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '541504', 'license': 'MIT', 'n_tokens_mistral': 982, 'n_tokens_neox': 930, 'n_words': 165, 'path': 'REINFORCE_old.ipynb', 'repo_name': 'vicgalle/optimal-reward-design', 'size': '2246'}