Distribution.sample() graph breaks

X3N4: <p>I’m working in Reinforcement Learning and was pretty excited about the potential benefits of torch.compile (jax.jit can speed up RL quite dramatically). I’ve noticed however that sampling from a torch probability distribution breaks the FX Graph. That’s a bummer for RL, where the policy usually parameterizes a probability distribution.</p>
<p>Here’s an example network:</p>
<pre><code class="lang-python">class Actor(nn.Module):
    def __init__(self, envs):
        super().__init__()
        obs_shape = envs.single_observation_space.shape
        self.conv = nn.Sequential(
            layer_init(nn.Conv2d(obs_shape[0], 32, kernel_size=8, stride=4)),
            nn.ReLU(),
            layer_init(nn.Conv2d(32, 64, kernel_size=4, stride=2)),
            nn.ReLU(),
            layer_init(nn.Conv2d(64, 64, kernel_size=3, stride=1)),
            nn.Flatten(),
        )

        with torch.inference_mode():
            output_dim = self.conv(torch.zeros(1, *obs_shape)).shape[1]

        self.fc1 = layer_init(nn.Linear(output_dim, 512))
        self.fc_logits = layer_init(nn.Linear(512, envs.single_action_space.n))

    def forward(self, x):
        x = F.relu(self.conv(x))
        x = F.relu(self.fc1(x))
        logits = self.fc_logits(x)

        return logits

    def get_action(self, x):
        logits = self(x / 255.0)
        policy_dist = Categorical(logits=logits)
        action = policy_dist.sample() # &lt;- GRAPH BREAK
        action_probs = policy_dist.probs
        log_prob = F.log_softmax(logits, dim=1)
        return action, log_prob, action_probs
</code></pre>
<p>Trying to compile the network or any kind of loss will always break the graph at the indicated line.</p>
<p>As I don’t understand torch.compile very well I was wondering:</p>
<ul>
<li>Will this be fixed in a future update?</li>
<li>If not: Is there are workaround I can use?</li>
</ul>
<p>Thanks in advance for the help and sorry if it’s a bit of a noob question.</p>

marksaroufim: <p>Hi <a class="mention" href="/u/x3n4">@X3N4</a> could you please share a full repro?</p>

KFrank: <p>Hi X!</p>
<aside class="quote no-group quote-modified" data-username="X3N4" data-post="1" data-topic="182680" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/x3n4/48/13899_2.png" class="avatar"> X3N4:</div>
<blockquote>
<pre data-code-wrap="python"><code class="lang-python">        policy_dist = Categorical(logits=logits)
        action = policy_dist.sample() # &lt;- GRAPH BREAK
</code></pre>
<p>Trying to compile the network or any kind of loss will always break the graph at the indicated line.</p>
</blockquote>
</aside>
<p>There are two thing going on here:  First, you can’t backpropagate through<br>
<code>.sample()</code>.  But even if you could, the result of <code>Categorical.sample()</code> is<br>
an integer, and you can’t backpropagate through integers (because they are<br>
discrete and therefore not usefully differentiable).</p>
<p>Depending on your specific use case, you may consider backpropagating<br>
through <code>.log_prob (action)</code> (where <code>action</code> is the result of <code>.sample()</code>),<br>
which does make sense (even for <code>Categorical</code>) and is permitted.</p>
<p>Quoting from the documentation for <a href="https://pytorch.org/docs/2.0/distributions.html" rel="noopener nofollow ugc">torch.distributions:</a></p>
<pre data-code-wrap="python"><code class="lang-python">probs = policy_network(state)
# Note that this is equivalent to what used to be called multinomial
m = Categorical(probs)
action = m.sample()
next_state, reward = env.step(action)
loss = -m.log_prob(action) * reward
loss.backward()
</code></pre>
<p>Such a scheme may (or may not) make sense for your use case.</p>
<p>Good luck!</p>
<p>K. Frank</p>

X3N4: <p>Hi KFrank,</p>
<p>thanks for the suggestion! Indeed you are correct that one has to backpropagate using the log-probabilities of the action, which is called the policy gradient. Your example is the REINFORCE algorithm.</p>
<p>However, please note that I do not intend to backpropagate through <code>Categorical.sample()</code>. Instead, I’m using a more elaborate version of the algorithm you gave in your example: SAC-discrete, an off-policy actor-critic algorithm where the actor (=policy) is updated similarly to your example.</p>

X3N4: <p>Hi <a class="mention" href="/u/marksaroufim">@marksaroufim</a></p>
<p>thank you very much for answering so quickly. I really appreciate it.</p>
<p>Here’s the link to a repro with a relevant example:</p><aside class="onebox allowlistedgeneric" data-onebox-src="https://github.com/timoklein/rl_graph_breaks">
  <header class="source">
      <img src="https://github.githubassets.com/favicons/favicon.svg" class="site-icon" width="32" height="32">

      <a href="https://github.com/timoklein/rl_graph_breaks" target="_blank" rel="noopener nofollow ugc">GitHub</a>
  </header>

  <article class="onebox-body">
    <div class="aspect-image" style="--aspect-ratio:690/345;"><img src="https://opengraph.githubassets.com/57950757311a30c9a48afec661816c31f146cf6d9debea15ece612e71cdf82a6/timoklein/rl_graph_breaks" class="thumbnail" width="690" height="345"></div>

<h3><a href="https://github.com/timoklein/rl_graph_breaks" target="_blank" rel="noopener nofollow ugc">GitHub - timoklein/rl_graph_breaks: An example of torch.compile graph breaks...</a></h3>

  <p>An example of torch.compile graph breaks in RL code using SAC-discrete as an example - GitHub - timoklein/rl_graph_breaks: An example of torch.compile graph breaks in RL code using SAC-discrete as ...</p>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>Contains the following:</p>
<ul>
<li>A simplified example using fake data with only torch dependencies.</li>
<li>A realistic example with the full algorithm I’m using (SAC-discrete).</li>
<li>
<code>dynamo.explain()</code> printouts in the README.</li>
</ul>
<p>I’ve tried using <code>no_grad</code> around the sampling, but that doesn’t work. An option might be to change the structure of the code, i.e. separate action generation (which uses <code>.sample()</code>) from calculating the log-probabilities.</p>
<p>I want to mention that the general structure of this example (sampling from a distribution and then using <code>log_prob</code> to backpropagate) is used widely in deep RL in many policy gradient/actor-critic algorithms (e.g. REINFORCE, A2C, PPO, SAC-discrete).</p>

X3N4: <aside class="quote no-group" data-username="X3N4" data-post="5" data-topic="182680">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/x3n4/48/13899_2.png" class="avatar"> X3N4:</div>
<blockquote>
<p>An option might be to change the structure of the code, i.e. separate action generation (which uses <code>.sample()</code>) from calculating the log-probabilities.</p>
</blockquote>
</aside>
<p>Just want to mention that this was able to fix the mentioned graph breaks.</p>