Trajectories.log_rewards lacks a detach guard: a differentiable state would silently change the TB estimator

[josephdviviano]

Summary

Trajectories.log_rewards computes the log-reward from terminating_states with no detach():

https://github.com/GFNOrg/torchgfn/blob/master/src/gfn/containers/trajectories.py#L250

if self._log_rewards is None:
    self._log_rewards = self.env.log_reward(self.terminating_states)

This is not a live bug — Sampler guards its output properly, so the states reaching this line are always grad-free today. I'm raising it as defence-in-depth, because it's the one place where a differentiable state would silently change the estimator rather than raise.

Why it's inert today

Two independent guards keep sampled states off the graph:

• src/gfn/samplers.py:88-89 and :228-229 — with torch.no_grad(): actions_tensor = dist.sample()
• src/gfn/utils/distributions.py:258-262 — IsotropicGaussian.sample is a literal reparameterisation (loc + scale * noise) and would be differentiable, but actions_detach defaults to True, and no call site in the package passes False.

I checked every release I could obtain (1.0.0 through 2.4.1, plus master) and the no_grad guard is continuous throughout. Nice.

What would happen if a differentiable state ever reached it

For continuous environments with a smooth reward, a gradient flows from log R(x_T) straight back into the policy parameters. Verified on 2.4.1:

import torch
from gfn.gym.diffusion_sampling import DiffusionSampling

env = DiffusionSampling(target_str="gmm2", target_kwargs={}, num_discretization_steps=10)
theta = torch.zeros(1, env.s0.shape[-1], requires_grad=True)
st = env.states_from_tensor(theta.detach().clone())
st.tensor = theta                      # a differentiable state
lr = env.log_reward(st)
print(lr.requires_grad, type(lr.grad_fn).__name__)
print(bool((torch.autograd.grad(lr.sum(), theta)[0] != 0).any()))

True LogsumexpBackward0
True

That matters because TB's residual is delta = log P_F - log P_B - log R(x_T). With grad-free states, theta enters only through log P_F and the gradient is the standard score-function/REINFORCE form. If x_T is differentiable in theta, a second pathwise term appears through log R, and the loss stops being the estimator it is documented to be — it silently becomes a hybrid that also rewards moving samples toward high-reward regions directly.

Scope: this only applies to environments whose reward is differentiable in the state. BoxPolar is unaffected — its reward is piecewise-constant ((0.25 < ax).prod(-1)), so it carries no gradient either way. The DiffusionSampling targets are the affected class.

How it could be reached

Not through Sampler, but:

1. Constructing Trajectories from states produced anywhere other than Sampler (custom rollouts, replay/relabelling, imported trajectories).
2. IsotropicGaussian(..., actions_detach=False) combined with any path that bypasses the sampler's no_grad.

In both cases it fails silently: training runs, loss decreases, and the gradient is simply a different estimator than intended. There's no assertion downstream that would catch it.

Suggested fix

A one-liner at the point of computation:

self._log_rewards = self.env.log_reward(self.terminating_states).detach()

If you'd rather fail loudly than silently correct — arguably better, since a differentiable state arriving here means something upstream is already not doing what its author expected:

assert not self.terminating_states.tensor.requires_grad, (
    "terminating states carry a grad_fn; log_rewards would introduce a pathwise "
    "gradient and silently change the TB estimator"
)

Happy to open a PR for either if useful.

Context

Found while auditing our own continuous-GFlowNet code, where the equivalent sampler did not detach and we'd been training with the hybrid estimator without realising. torchgfn gets this right; I wanted to flag the one remaining unguarded spot in case it's cheap to close.

Tested against torchgfn 2.4.1 (PyPI sdist) and cross-checked against master @ 654b828.