Hi albanD,
Thanks for you answer, sorry for my late response (did not get a notification).
The second solution works, but now how do I compute the second derivative?

for i in range(200):

  y = get_expected_return(policy)
  grad = torch.autograd.grad(y, policy.parameters(), create_graph=True)

  for gvec, (name,theta) in zip(grad, policy.named_parameters()):
    gvec=gvec.flatten()
    
    H = [torch.autograd.grad(g, theta, create_graph=True)[0].unsqueeze(0) for g in gvec]

Is this correct?
Consider two parameter sets, x1 and x2. The H should be computed by:

In my code, where am I computing right upper and left lower elements?