Hessian Vector Product for discounted-return-based loss function

username: DINHQuangDung1999
cooked: <p>Hi all,</p>
<p>I am trying to implement a Hessian-guided SGD optimizer for a RL task, which requires the hessian vector product computed in an online manner. I have tried computing HVP using the following two approaches.</p>
<p><strong>1st approach.</strong></p>
<pre><code class="lang-auto">params_with_grad = []
for name, params in self.actor.named_parameters():
    if '_module' in name:
        params_with_grad.append(params)
actor_loss_PW = self.compute_actor_loss(pathwise_grad=True)
        
PW_grads = torch.autograd.grad(actor_loss_PW, params_with_grad,create_graph=True)
num_param = sum(p.numel() for p in params_with_grad)
v = [torch.rand(PW_grads[i].shape).to('cuda') for i in range(len(PW_grads))]
HVP = torch.autograd.grad(PW_grads, params_with_grad, v, retain_graph=True)
</code></pre>
<p><strong>2nd approach.</strong></p>
<pre><code class="lang-auto">params_with_grad = []
for name, params in self.actor.named_parameters():
    if '_module' in name:
        params_with_grad.append(params)
actor_loss_PW = self.compute_actor_loss(pathwise_grad=True)
        
PW_grads = torch.autograd.grad(actor_loss_PW, params_with_grad,create_graph=True)
PW_grads = torch.cat([e.flatten() for e in PW_grads]) # flatten
num_param = sum(p.numel() for p in params_with_grad)
v = torch.rand(num_param).to('cuda')
HVP = torch.autograd.grad(PW_grads, params_with_grad, v, retain_graph=True)
</code></pre>
<p>Both approaches ended up with</p>
<blockquote>
<p>*** IndexError: pop from empty list.</p>
</blockquote>
<p>The weird thing is there is no further explanation for this error. I wrapped some layers of my actor network with GradSampleModule from Opacus.</p>
<p><strong>UPDATE.</strong> I found that removing the GradSampleModule wrapper solves the problem.  However, I need GradSampleModule to get individual gradient samples.</p>
<p>The calculation of the loss is as follow</p>
<pre><code class="lang-auto">    def compute_actor_loss(self, update_metrics = True, pathwise_grad = True):
        rew_acc = torch.zeros(self.num_envs, dtype=torch.float32, device=self.device)
        gamma = torch.ones(self.num_envs, dtype=torch.float32, device=self.device)
        if not pathwise_grad:
            logprob_acc = torch.zeros(self.num_envs, dtype=torch.float32, device=self.device)
        with torch.no_grad():
            if self.obs_rms is not None:
                obs_rms = deepcopy(self.obs_rms)
        # actor_loss = torch.zeros(self.num_envs, dtype=torch.float32, device=self.device)
        
        # initialize trajectory to cut off gradients between episodes.
        obs = self.env.initialize_trajectory()
        obs = self._convert_obs(obs)
        # breakpoint()
        if self.obs_rms is not None:
            # update obs rms
            with torch.no_grad():
                for k, v in obs.items():
                    self.obs_rms[k].update(v)
            # normalize the current obs
            obs = {k: obs_rms[k].normalize(v) for k, v in obs.items()}

        # collect trajectories and compute actor loss
        for i in range(self.horizon_len):
            # take env step
            if pathwise_grad:
                actions = self.get_actions(obs, eval = False)
            else:
                actions, logprob = self.get_actions(obs, eval = False, pathwise = False)
            obs, rew, done, extra_info = self.env.step(actions)
            obs = self._convert_obs(obs)

            with torch.no_grad():
                raw_rew = rew.clone()
            # scale the reward
            rew = self.reward_shaper(rew)

            if self.obs_rms is not None:
                # update obs rms
                with torch.no_grad():
                    for k, v in obs.items():
                        self.obs_rms[k].update(v)
                # normalize the current obs
                obs = {k: obs_rms[k].normalize(v) for k, v in obs.items()}

            done_env_ids = done.nonzero(as_tuple=False).squeeze(-1)
            not_done_env_ids = (done == 0).nonzero(as_tuple=False).squeeze(-1)

            # compute actor loss
            rew_acc[not_done_env_ids] = rew_acc[not_done_env_ids] + gamma[not_done_env_ids] * rew[not_done_env_ids]
            if not pathwise_grad:
                logprob_acc[not_done_env_ids] += logprob[not_done_env_ids].sum(dim = 1)

            # compute gamma for next step
            gamma = gamma * self.gamma

            # clear up gamma and rew_acc for done envs
            gamma[done_env_ids] = 1.0

        if pathwise_grad:
            actor_loss = -rew_acc
        else:
            _loss = -rew_acc
            _logprob = logprob_acc
            baseline = _loss.mean()
            adv = (_loss - baseline).detach()
            actor_loss = adv * _logprob
                    
        actor_loss /= self.horizon_len * self.num_envs
        actor_loss = actor_loss.sum()
        return actor_loss
</code></pre>

username: valerian.rey
cooked: <p>You can get per-sample gradients using the <code>torch.func</code> API (see <a href="https://docs.pytorch.org/tutorials/intermediate/per_sample_grads.html" rel="noopener nofollow ugc">this</a>), and you can use the <code>torch.autograd.functional</code> API to compute <a href="https://docs.pytorch.org/docs/stable/generated/torch.autograd.functional.hvp.html" rel="noopener nofollow ugc">hvp</a> (or <a href="https://docs.pytorch.org/docs/stable/generated/torch.autograd.functional.vhp.html" rel="noopener nofollow ugc">vhp</a> - apparently it’s more efficient). So maybe you can try to combine both, instead of relying on Opacus?</p>
<p>Could you also clarify what you want to do, mathematically? In particular:</p>
<ol>
<li>What vector do you want to use for the hvp? (I’m guessing that the <code>torch.rand</code> in the code you shared is just for the sake of the example).</li>
<li>What do you want to do with the per-sample gradients?</li>
</ol>

username: DINHQuangDung1999
cooked: <p>I am trying to compute a hybrid between the score function and pathwise gradient. I sampled one where the reward is detached and gradient only flows through the logprob, store the grad, empty the gradient (for SF grad), then repeat for the PW grad. I wanted to compute the sample variance of each grad - which requires knowing the sample gradients of each trajectory - leading me to using Opacus.</p>
<p>I found that the <code>torch.func.grad/vmap</code>only works for static dataset (e.g. in supervised/unsupervised setting), whereas my interest is at continuous control/RL. Specifically, this approach requires knowing the networks params and the datasets as inputs,  for <code>torch.func.grad</code>, so I did not manage to use it for my interest. If you know how one can do this, I would love to hear it!</p>
<p>In terms of ideaology, I was trying to look at the directional curvature of the loss in the direction of the current gradient, so <code>v</code> would be the current gradient tensor, rather than a random tensor.</p>
<p><strong>UPDATE.</strong> I found that I can just still wrap my layers in <code>GradSampleModule</code> and then manually enable/disable the hook by Opacus to overcome this issue.</p>
<pre><code class="lang-auto">for i in range(epochs) do:
1. self.actor_mlp.layer_names.enable_hooks()
2. compute and store SF grad samples
3. compute and store PW grad samples
4. self.actor_mlp.layer_names.disable_hooks()
5. compute hvp and directional curvature 
6. update net params
end for
</code></pre>
<p>Perhaps the reason is Opacus’s hooks not functioning properly given multiple backward pass (2 in my case). I believe this problem has been discussed somewhere on this forum.</p>