Stability of a 2-pass update through a shared layer?

mike_and_ike: <p>Hey all! Say I’ve got a network like the following</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/a/6/a65b69d7ed3d7f6d1ab0725390c751baa05bb3df.png" alt="2021-10-30-175114_1920x1080_scrot" data-base62-sha1="nJF6YujeT380cNVYTl9XXC5OUYT" width="583" height="285"></p>
<p>Training the RL algorithm, DDPG, comprises 2 steps:</p>
<p><a href="https://spinningup.openai.com/en/latest/_images/math/5811066e89799e65be299ec407846103fcf1f746.svg" rel="noopener nofollow ugc">OpenAI Algorithm Outline</a></p>
<p>I see a number of implementations online, but most do not have a shared layer so I don’t have a baseline to compare against.</p>
<p><strong>Line 12 and 13</strong></p>
<p>This portion describes how to compute the loss and gradient update through the first head (Q-&gt;action-value)</p>
<p><strong>Line 14</strong></p>
<p>This portion describes how to compute the loss and gradient update through the second head, Policy. I’ve coded it such that on every pass of the <code>actor-critic</code> I receive 2 outputs, the policy and the Q-value.</p>
<p>I’ve separated this into 2 optimization steps:</p>
<h1><a name="step-1-1" class="anchor" href="#step-1-1"></a>Step 1</h1>
<pre><code class="lang-auto">self.actor_critic.critic_optimizer.zero_grad()
 _, y_hat_critic = self.actor_critic(s_batch, a_batch)
critic_loss = F.mse_loss(y_hat_critic, y_critic)
critic_loss.backward()
self.custom_gradient_scaling_for_shared(self.actor_critic.parameters())
self.actor_critic.critic_optimizer.step()
</code></pre>
<h1><a name="step-2-2" class="anchor" href="#step-2-2"></a>Step 2</h1>
<pre><code class="lang-auto">self.actor_critic.actor_optimizer.zero_grad()
 _, y_hat_policy = self.actor_critic(s_batch)
policy_loss = T.mean(-y_hat_policy)
policy_loss.backward()
self.custom_gradient_scaling_for_shared(self.actor_critic.parameters())
self.actor_critic.actor_optimizer.step()
</code></pre>
<p>P.s I’ve verified that in the first case, the policy head’s gradient update is zero, and in the second, although a gradient is generated for the Q-value head, it is not applied.</p>
<p>However, I am facing an issue in terms of determining if the 2 pass update would cause the shared layer to run into stability issues. I’ve got <code>self.custom_gradient_scaling_for_shared(self.actor_critic.parameters())</code> which determines if the update is for the shared layer, and if it is, it just scales the gradient.</p>
<p>Does anyone have insights into this?</p>

tom: <p>I don’t think it necessarily is a problem, but if you disregard the update between step 1 and step 2 and the potentially different optimizer hyperparameters, you are effectively optimizing the shared weights to a sum of the two objective functions. Depending on the size relations of the two, this may make more or less sense.</p>
<p>Best regards</p>
<p>Thomas</p>
<p>P.S.: Is it intentional that you take the second output of actor_critic for both?</p>

mike_and_ike: <blockquote>
<p>I don’t think it necessarily is a problem, but if you disregard the update between step 1 and step 2 and the potentially different optimizer hyperparameters, you are effectively optimizing the shared weights to a sum of the two objective functions</p>
</blockquote>
<p>True, true! I was just wondering if this might pose a problem since, as we both mentioned, it’s optimizing over two objectives (hence my scaling the gradient update for the shared layers).</p>
<blockquote>
<p>P.S.: Is it intentional that you take the second output of actor_critic for both?</p>
</blockquote>
<p>Yeah, the first output is the output of the actor, and the second is the output of the critic. Our <code>actor_critic</code> has 2 “cases”:</p>
<ol>
<li><code>self.actor_critic(state, action)</code>, the critic head uses the passed-in action to output the Q-value. (line 13, the LHS of the optimization eqn)</li>
</ol>
<p>Q_{\phi}(s, a) where (s, a) are sampled from the replay buffer.</p>
<ol start="2">
<li><code>self.actor_critic(state)</code> where we do not pass in an action. Here, we use our policy head to output a policy and we then output the Q-value from that (Note how in line 14, of the algorithm, we are taking the derivative of the Q-value but WRT the actor’s parameters)</li>
</ol>
<p>Q_{\phi}(s, mu_{\theta}(a)). I needed to be careful with how I generated the gradient but I think I did it correctly.</p>
<p>The below is the code implementing “case 2”</p>
<p><strong>Optimizer</strong></p>
<pre data-code-wrap="python"><code class="lang-python">params = list(self.shared_module.parameters()) + list(self.mu.parameters())
self.actor_optimizer = optim.Adam(params, lr=self.actor_lr)
</code></pre>
<p>Forward pass in “case 2”</p>
<pre><code class="lang-auto">compressed: T.Tensor = self.shared_layer(state)
########################################
# Actor runs first
########################################
mu_val = self.mu(F.relu(compressed))
actions: T.Tensor = T.tanh(mu_val)
leq, geq = actions &lt;= 1, actions &gt;= -1
if not torch.all(leq * geq):
    raise Exception("Our actions were out of the expected range")

########################################
# Critic runs now
########################################
action_value = F.relu(self.action_value(actions))
state_action_value = self.q(
    F.relu(T.add(compressed, action_value))
)
return actions, state_action_value
</code></pre>
<p>And the update step is as in my original post.</p>
<hr>
<p>Does anything look egregiously wrong?</p>

mike_and_ike: <p><a class="mention" href="/u/tom">@tom</a></p>
<p>Sorry to ping you directly, but did you have any follow up thoughts? If not I’ll just close this question since it seems like there’s not much activity</p>