Extend modules to support more kinds of diffusion processes

[josephdviviano]

FYI: a natural extension would be using a different diffusion process (pinned Brownian motion -> Ornstein-Uhlenbeck process), which requires defining new forward and backward Estimators, while we can share the modules and the IsotropicGaussian distribution.

Originally posted by @hyeok9855 in #419 (comment)

[josephdviviano]

@claude can you formulate a solution to this that focuses on elegance and compositionally of the underlying classes- the code should be compact and expressive to allow the user to produce a large number of different diffusion processes with minimal code written. We will likely need to iterate a few times to get the design down, so feel free to ask many questions or propose different ideas for us to debate .

[claude]

Claude finished @josephdviviano's task in 2m 54s —— View job

---

Design Proposal: Composable Diffusion Processes

• Read CLAUDE.md and understand repo conventions
• Explore current diffusion process implementation (modules, estimators, distributions)
• Understand the PR Diffusion Sampling #419 context
• Formulate design proposal with composability focus
• Post design questions/options for discussion

---

Analysis of the Current Design

The current code has a clear issue: the SDE math is hardcoded inside the estimator classes. PinnedBrownianMotionForward.to_probability_distribution (lines 1439-1521 of estimators.py) directly embeds the Brownian bridge math, and any new process requires copy-pasting this entire estimator. There's even a # TODO: support OU process comment on line 1391.

The modules (DiffusionPISGradNetForward, DiffusionFixedBackwardModule, etc.) are already process-agnostic — they just output drift vectors. So the key separation to make is: pull the SDE math out of the estimator and into its own abstraction.

---

Proposed Design: SDE Strategy Object

The core idea is a single abstract SDE class that captures three primitives needed to define any Itô diffusion for GFlowNet training:

class SDE(ABC):
    """Mathematical definition of a stochastic differential process.
    
    An SDE defines the forward reference drift, diffusion noise, and the
    backward bridge kernel (used to compute P_B analytically).
    """

    @abstractmethod
    def reference_drift(self, x: Tensor, t: Tensor) -> Tensor:
        """The intrinsic (uncontrolled) forward drift, added to the NN output."""

    @abstractmethod
    def diffusion_std(self, t: Tensor, dt: float) -> Tensor:
        """Isotropic noise std for a single step of size dt."""

    @abstractmethod
    def bridge_mean_std(self, x: Tensor, t: Tensor, dt: float) -> tuple[Tensor, Tensor]:
        """Analytic backward bridge kernel: (base_mean, base_std) for P_B."""

Two concrete implementations cover the immediate need:

class PinnedBrownianMotion(SDE):
    """Reference SDE: dx = u(x,t)dt + σdW, pinned to 0 at t=0."""
    def __init__(self, sigma: float): ...
    def reference_drift(self, x, t): return torch.zeros_like(x)
    def diffusion_std(self, t, dt): return self.sigma * dt**0.5
    def bridge_mean_std(self, x, t, dt):
        t_safe = t.clamp(min=dt * 1e-4)
        mean = x * dt / t_safe
        std = self.sigma * (dt * (t_safe - dt).clamp(min=0) / t_safe).sqrt()
        return mean, std

class OrnsteinUhlenbeck(SDE):
    """Reference SDE: dx = θ(μ - x)dt + σdW."""
    def __init__(self, sigma: float, theta: float, mu: float = 0.0): ...
    def reference_drift(self, x, t): return self.theta * (self.mu - x)
    def diffusion_std(self, t, dt): return self.sigma * dt**0.5
    def bridge_mean_std(self, x, t, dt):
        # OU bridge formula — analytically tractable, see Doob's h-transform
        ...

Then two unified estimators replace the four current classes:

class DiffusionForwardEstimator(PolicyMixin, Estimator):
    """Forward policy for any SDE.  Adds reference drift to NN control."""
    def __init__(self, s_dim, module, sde: SDE, dt, n_variance_outputs=0): ...
    
    def to_probability_distribution(self, states, module_output, **kwargs):
        x, t = states.tensor[:, :-1], states.tensor[:, [-1]]
        drift = self.sde.reference_drift(x, t) + module_output[:, :self.s_dim]
        std = self.sde.diffusion_std(t, self.dt)
        return IsotropicGaussian(self.dt * drift, std)

class DiffusionBackwardEstimator(PolicyMixin, Estimator):
    """Backward policy for any SDE.  Corrects the analytic bridge kernel."""
    def __init__(self, s_dim, module, sde: SDE, dt, pb_scale_range=0.1): ...
    
    def to_probability_distribution(self, states, module_output, **kwargs):
        x, t = states.tensor[:, :-1], states.tensor[:, [-1]]
        base_mean, base_std = self.sde.bridge_mean_std(x, t, self.dt)
        correction = module_output[:, :self.s_dim] * self.pb_scale_range
        return IsotropicGaussian(base_mean + correction, base_std)

What this enables at call sites

# Today (PBM) — identical behavior to current code
sde = PinnedBrownianMotion(sigma=1.0)
pf = DiffusionForwardEstimator(s_dim=2, module=pf_net, sde=sde, dt=dt)
pb = DiffusionBackwardEstimator(s_dim=2, module=pb_net, sde=sde, dt=dt)

# Tomorrow (OU) — single line change
sde = OrnsteinUhlenbeck(sigma=1.0, theta=2.0)
pf = DiffusionForwardEstimator(s_dim=2, module=pf_net, sde=sde, dt=dt)
pb = DiffusionBackwardEstimator(s_dim=2, module=pb_net, sde=sde, dt=dt)

# Custom (user-defined) — zero changes to library code
class MySDE(SDE):
    def reference_drift(self, x, t): ...
    def diffusion_std(self, t, dt): ...
    def bridge_mean_std(self, x, t, dt): ...

---

Design Questions to Debate

Q1 — Where does the exit action logic live?
Currently PinnedBrownianMotionForward.to_probability_distribution checks if the next step reaches terminal time and replaces the output with -inf. This is a GFlowNet bookkeeping concern, not an SDE concern. My instinct is it belongs in DiffusionForwardEstimator (process-agnostic), but it needs dt from the estimator and DIFFUSION_TERMINAL_TIME_EPS. Does that seem right, or should it live in the environment's step()?

Q2 — Learned variance: estimator or SDE?
n_variance_outputs controls whether the NN outputs extra dims that are interpreted as learned variance corrections. This is a learning/parameterization choice, not a property of the underlying SDE. I'd keep it on the estimator. But the current design has DiffusionPISGradNetForward.learn_variance=True setting out_dim = s_dim + 1 — the module and estimator are tightly coupled. Should we keep this or break it apart?

Q3 — Single DiffusionEstimator(is_backward=True/False) vs. two classes?
The forward and backward estimators share structure but the math in to_probability_distribution differs enough that two classes feel cleaner to read. Counter-argument: a single class + is_backward flag matches the existing pattern (DiscretePolicyEstimator uses a flag). What do you prefer?

Q4 — reference_drift in the forward estimator: additive to NN, or blended differently?
For PBM the reference drift is 0, so this is a no-op. For OU the total drift is θ(μ-x) + u(x,t) where u is the NN output. This is the natural parameterization for diffusion GFlowNets (the NN learns to steer away from the reference). But one could also parameterize where the NN outputs the full drift (including OU terms), which changes the learning problem. Any preference?

Q5 — IsotropicGaussian vs. richer distributions?
The issue title mentions "more kinds of diffusion processes". Some SDEs have anisotropic diffusion or state-dependent noise. Should we generalize IsotropicGaussian now (e.g. DiagonalGaussian) or keep it isotropic and note it as a follow-up?

Q6 — Backward compatibility strategy?
Option A: Deprecate PinnedBrownianMotionForward/Backward, replace with new classes.
Option B: Make them subclasses of the new unified estimators (PinnedBrownianMotionForward(DiffusionForwardEstimator) with a convenience __init__).
Option B is least disruptive and keeps existing tutorials working.

---

Let me know which questions are most important to settle first and I'll start sketching the implementation. The OU bridge_mean_std formula is the one piece I want to double-check mathematically before committing to it — want to share a reference or confirm the Doob's h-transform approach is the intended derivation?