How to replace usage of "retain_graph=True"

username: Yuerno
cooked: <p>Hi all. I’ve generally seen it recommended against using the retain_graph parameter, but I can’t seem to get a piece of my code working without it. In particular, this is the code where I use retain_graph (the goal here is to optimize a single scalar <code>cql_alpha</code> value:</p>
<pre><code class="lang-auto"># calculate eq. 5 in updated SAC paper
        qf1_values = qf1(observations)
        qf2_values = qf2(observations)
        qf1_a_values = qf1_values.gather(1, actions)
        qf2_a_values = qf2_values.gather(1, actions)
        qf1_loss = F.mse_loss(qf1_a_values, next_q_value)
        qf2_loss = F.mse_loss(qf2_a_values, next_q_value)
        qf_loss = 0.5 * (qf1_loss + qf2_loss)  # scaling added from CQL

        # calculate CQL regularization loss
        cql_qf1_diff = torch.logsumexp(qf1_values, dim=1).mean() - qf1_a_values.mean()
        cql_qf2_diff = torch.logsumexp(qf2_values, dim=1).mean() - qf2_a_values.mean()
        if args.cql_autotune:
            cql_alpha = torch.clamp(torch.exp(cql_log_alpha), min=0.0, max=1000000.0)
            cql_qf1_loss = cql_alpha * (cql_qf1_diff - args.difference_threshold)
            cql_qf2_loss = cql_alpha * (cql_qf2_diff - args.difference_threshold)
            cql_alpha_loss = -(cql_qf1_loss + cql_qf2_loss)

            # ---------- update cql_alpha ---------- #
            cql_a_optimizer.zero_grad()
            cql_alpha_loss.backward(retain_graph=True)
            cql_a_optimizer.step()
        else:
            cql_qf1_loss = cql_alpha * cql_qf1_diff
            cql_qf2_loss = cql_alpha * cql_qf2_diff

        # calculate final q-function loss which is a combination of Bellman error and CQL regularization
        cql_qf_loss = cql_qf1_loss + cql_qf2_loss
        qf_loss = cql_qf_loss + qf_loss

        # calculate eq. 6 in updated SAC paper
        q_optimizer.zero_grad()
        qf_loss.backward()
        q_optimizer.step()
</code></pre>
<p>I tried the following approach to remove the usage, but while it seems to work initially, eventually the Q-values explode to infinity with this approach.</p>
<pre><code class="lang-auto"># calculate eq. 5 in updated SAC paper
        qf1_values = qf1(observations)
        qf2_values = qf2(observations)
        qf1_a_values = qf1_values.gather(1, actions)
        qf2_a_values = qf2_values.gather(1, actions)
        qf1_loss = F.mse_loss(qf1_a_values, next_q_value)
        qf2_loss = F.mse_loss(qf2_a_values, next_q_value)
        qf_loss = 0.5 * (qf1_loss + qf2_loss)  # scaling added from CQL

        # calculate CQL regularization loss
        if args.cql_autotune:
            cql_alpha = torch.clamp(
                cql_log_alpha.exp().detach(), min=0.0, max=1000000.0
            )
            cql_qf1_loss = (
                torch.logsumexp(qf1_values, dim=1).mean() - qf1_a_values.mean()
            )
            cql_qf2_loss = (
                torch.logsumexp(qf2_values, dim=1).mean() - qf2_a_values.mean()
            )
            cql_qf1_loss = cql_alpha * (cql_qf1_loss - args.difference_threshold)
            cql_qf2_loss = cql_alpha * (cql_qf2_loss - args.difference_threshold)
        else:
            cql_qf1_loss = cql_alpha * (
                torch.logsumexp(qf1_values, dim=1).mean() - qf1_a_values.mean()
            )
            cql_qf2_loss = cql_alpha * (
                torch.logsumexp(qf2_values, dim=1).mean() - qf2_a_values.mean()
            )

        # calculate final q-function loss which is a combination of Bellman error and CQL regularization
        cql_qf_loss = cql_qf1_loss + cql_qf2_loss
        qf_loss = cql_qf_loss + qf_loss

        # calculate eq. 6 in updated SAC paper
        q_optimizer.zero_grad()
        qf_loss.backward()
        q_optimizer.step()

        # ---------- update cql_alpha ---------- #
        if args.cql_autotune:
            with torch.no_grad():
                cql_qf1_diff_loss = cql_qf1_loss - args.difference_threshold
                cql_qf2_diff_loss = cql_qf2_loss - args.difference_threshold

            cql_alpha = torch.clamp(cql_log_alpha.exp(), min=0.0, max=1000000.0)
            cql_alpha_loss = cql_alpha * -(cql_qf1_diff_loss + cql_qf2_diff_loss)

            cql_a_optimizer.zero_grad()
            cql_alpha_loss.backward()
            cql_a_optimizer.step()
</code></pre>
<p>Would anyone be able to help me figure out what I might be doing wrong?</p>

username: vmoens
cooked: <p>I personally think it’s ok to use <code>retain_graph</code> if you have multiple losses where part of the graph is shared. You should be careful about calling zero_grad where appropriate to make sure that one backward does not impact another (you generally don’t want gradients to accumulate in multiple places).</p>
<p>To answer your question, I don’t see any obvious difference between the two. It feels like your first example could run without <code>retain_graph</code> as you only need to differentiate through a very small graph.<br>
The second solution definitely makes more sense.</p>
<p>Exploding Q-values can happen for multiple reasons with SAC-like algos: wrong implementations, wrong optimization of the entropy coefficient, …<br>
One thing is also numerical instability of action sampled around the bounds of the action space when using transforms. You may want to look into that.<br>
One thing I’ve found to help is to use a small weight decay in the optimizer(s).</p>