From c5b064af4a5deb954fa0c49d2f0baed8b4a6d51d Mon Sep 17 00:00:00 2001 From: Joseph Viviano Date: Thu, 2 Apr 2026 22:38:24 -0400 Subject: [PATCH 01/10] Add optional custom log_rewards to all GFlowNet losses Add a `log_rewards` parameter to all loss/score methods across TB, DB, SubTB, FM, RTB, LPV, and their variants. When None (default), behavior is unchanged. When provided, the custom rewards are used instead of environment rewards, enabling intrinsic reward schemes. Based on the proposal in #312 by @Idriss-Malek, re-implemented on current master to resolve conflicts with the refactored codebase. Co-Authored-By: Claude Opus 4.6 (1M context) --- src/gfn/gflownet/base.py | 14 ++++- src/gfn/gflownet/detailed_balance.py | 18 +++++- src/gfn/gflownet/flow_matching.py | 14 ++++- src/gfn/gflownet/sub_trajectory_balance.py | 53 ++++++++++++----- src/gfn/gflownet/trajectory_balance.py | 68 +++++++++++++++++++--- tutorials/examples/train_hypergrid_gafn.py | 9 ++- 6 files changed, 147 insertions(+), 29 deletions(-) diff --git a/src/gfn/gflownet/base.py b/src/gfn/gflownet/base.py index 116315eb..3ac0d34f 100644 --- a/src/gfn/gflownet/base.py +++ b/src/gfn/gflownet/base.py @@ -420,6 +420,7 @@ def logz_parameters(self) -> list[torch.Tensor]: def get_scores( self, trajectories: Trajectories, + log_rewards: torch.Tensor | None = None, recalculate_all_logprobs: bool = True, env: Env | None = None, ) -> torch.Tensor: @@ -430,6 +431,11 @@ def get_scores( Args: trajectories: The Trajectories object to evaluate. + log_rewards: Optional custom log rewards tensor of shape + (n_trajectories,). When None, uses the environment rewards + from the trajectories. Useful for intrinsic rewards (see + "Towards Improving Exploration through Sibling Augmented + GFlowNets", Madan et al., ICLR 2025). recalculate_all_logprobs: Whether to re-evaluate all logprobs. env: The environment (unused in base TB, but required by some subclasses such as RTB and SubTB). @@ -448,11 +454,13 @@ def get_scores( total_log_pf_trajectories = log_pf_trajectories.sum(dim=0) # [N] total_log_pb_trajectories = log_pb_trajectories.sum(dim=0) # [N] - # cast: log_rewards is always set for terminating trajectories; - # assert is behind debug to avoid graph breaks in torch.compile. - log_rewards = cast(torch.Tensor, trajectories.log_rewards) + # Use custom log_rewards if provided, otherwise fall back to the + # trajectory's environment rewards. + if log_rewards is None: + log_rewards = cast(torch.Tensor, trajectories.log_rewards) if self.debug: assert log_rewards is not None + assert log_rewards.shape == (trajectories.batch_size,) # Fast path: skip clamp when log_reward_clip_min is disabled (default). if math.isfinite(self.log_reward_clip_min): log_rewards = log_rewards.clamp_min(self.log_reward_clip_min) diff --git a/src/gfn/gflownet/detailed_balance.py b/src/gfn/gflownet/detailed_balance.py index f09923df..2b43638e 100644 --- a/src/gfn/gflownet/detailed_balance.py +++ b/src/gfn/gflownet/detailed_balance.py @@ -209,6 +209,7 @@ def get_scores( self, env: Env, transitions: Transitions, + log_rewards: torch.Tensor | None = None, recalculate_all_logprobs: bool = True, ) -> torch.Tensor: r"""Calculates the scores for a batch of transitions. @@ -219,6 +220,11 @@ def get_scores( Args: env: The environment where the transitions are sampled from. transitions: The Transitions object to evaluate. + log_rewards: Optional custom log rewards tensor of shape + (n_transitions,). When None, uses the environment rewards + from the transitions. Useful for intrinsic rewards (see + "Towards Improving Exploration through Sibling Augmented + GFlowNets", Madan et al., ICLR 2025). recalculate_all_logprobs: Whether to re-evaluate all logprobs. Returns: @@ -261,11 +267,13 @@ def get_scores( is_terminating = transitions.is_terminating is_intermediate = ~is_terminating - # cast: log_rewards is always populated for valid transitions; - # assert is behind debug to avoid graph breaks in torch.compile. - log_rewards = cast(torch.Tensor, transitions.log_rewards) + # Use custom log_rewards if provided, otherwise fall back to the + # transition's environment rewards. + if log_rewards is None: + log_rewards = cast(torch.Tensor, transitions.log_rewards) if self.debug: assert log_rewards is not None + assert log_rewards.shape == (transitions.n_transitions,) if math.isfinite(self.log_reward_clip_min): log_rewards = log_rewards.clamp_min(self.log_reward_clip_min) log_F_s_next[is_terminating] = log_rewards[is_terminating] @@ -332,6 +340,7 @@ def loss( self, env: Env, transitions: Transitions, + log_rewards: torch.Tensor | None = None, recalculate_all_logprobs: bool = True, reduction: str = "mean", ) -> torch.Tensor: @@ -343,6 +352,8 @@ def loss( Args: env: The environment where the transitions are sampled from. transitions: The Transitions object to compute the loss with. + log_rewards: Optional custom log rewards tensor of shape + (n_transitions,). When None, uses the environment rewards. recalculate_all_logprobs: Whether to re-evaluate all logprobs. reduction: The reduction method to use ('mean', 'sum', or 'none'). Run with self.debug=False for improved performance. @@ -356,6 +367,7 @@ def loss( scores = self.get_scores( env, transitions, + log_rewards=log_rewards, recalculate_all_logprobs=recalculate_all_logprobs, ) scores = self.loss_fn(scores) diff --git a/src/gfn/gflownet/flow_matching.py b/src/gfn/gflownet/flow_matching.py index 10e7336d..24edec27 100644 --- a/src/gfn/gflownet/flow_matching.py +++ b/src/gfn/gflownet/flow_matching.py @@ -279,6 +279,7 @@ def loss( self, env: DiscreteEnv, states_container: StatesContainer[DiscreteStates], + log_rewards: torch.Tensor | None = None, recalculate_all_logprobs: bool = True, reduction: str = "mean", ) -> torch.Tensor: @@ -293,6 +294,11 @@ def loss( env: The discrete environment where the states are sampled from. states_container: The StatesContainer object containing both intermediary and terminating states. + log_rewards: Optional custom log rewards tensor of shape + (n_terminating_states,). When None, uses the environment rewards + from the states container. Useful for intrinsic rewards (see + "Towards Improving Exploration through Sibling Augmented + GFlowNets", Madan et al., ICLR 2025). recalculate_all_logprobs: Whether to re-evaluate all logprobs (unused for FM). reduction: The reduction method to use ('mean', 'sum', or 'none'). @@ -319,10 +325,16 @@ def loss( states_container.intermediary_states, reduction=reduction, ) + + terminating_log_rewards = ( + log_rewards + if log_rewards is not None + else states_container.terminating_log_rewards + ) rm_loss = self.reward_matching_loss( env, states_container.terminating_states, - states_container.terminating_log_rewards, + terminating_log_rewards, reduction=reduction, ) return fm_loss + self.alpha * rm_loss diff --git a/src/gfn/gflownet/sub_trajectory_balance.py b/src/gfn/gflownet/sub_trajectory_balance.py index 1edc84cd..eab226ab 100644 --- a/src/gfn/gflownet/sub_trajectory_balance.py +++ b/src/gfn/gflownet/sub_trajectory_balance.py @@ -208,6 +208,7 @@ def calculate_targets( is_terminal_mask: MaskTensor, sink_states_mask: MaskTensor, i: int, + log_rewards: torch.Tensor | None = None, ) -> TargetsTensor: """Calculates the targets tensor for the current sub-trajectory length. @@ -225,19 +226,26 @@ def calculate_targets( sink_states_mask: A mask of shape (max_length, batch_size) indicating whether the state is a sink state. i: The sub-trajectory length. + log_rewards: Optional custom log rewards tensor of shape + (n_trajectories,). When None, uses the environment rewards. + Useful for intrinsic rewards (see + "Towards Improving Exploration through Sibling Augmented + GFlowNets", Madan et al., ICLR 2025). Returns: The targets tensor of shape (max_length + 1 - i, batch_size). """ targets = torch.full_like(preds, fill_value=-float("inf")) - # Guard behind debug: log_rewards is always set for terminating trajectories; - # bare assert would cause a graph break in torch.compile. + if log_rewards is None: + # Guard behind debug: log_rewards is always set for terminating trajectories; + # bare assert would cause a graph break in torch.compile. + if self.debug: + assert trajectories.log_rewards is not None + # cast: log_rewards is always set for terminating trajectories. + log_rewards = cast(torch.Tensor, trajectories.log_rewards) if self.debug: - assert trajectories.log_rewards is not None - # cast: log_rewards is always set for terminating trajectories. - log_rewards = cast(torch.Tensor, trajectories.log_rewards)[ - trajectories.terminating_idx >= i - ] + assert log_rewards.shape == (trajectories.batch_size,) + log_rewards = log_rewards[trajectories.terminating_idx >= i] if math.isfinite(self.log_reward_clip_min): log_rewards = log_rewards.clamp_min(self.log_reward_clip_min) @@ -333,6 +341,7 @@ def calculate_masks( def get_scores( self, trajectories: Trajectories, + log_rewards: torch.Tensor | None = None, recalculate_all_logprobs: bool = True, env: Env | None = None, ) -> Tuple[List[torch.Tensor], List[torch.Tensor]]: @@ -341,6 +350,11 @@ def get_scores( Args: env: The environment where the trajectories are sampled from. trajectories: The batch of trajectories to evaluate. + log_rewards: Optional custom log rewards tensor of shape + (n_trajectories,). When None, uses the environment rewards. + Useful for intrinsic rewards (see + "Towards Improving Exploration through Sibling Augmented + GFlowNets", Madan et al., ICLR 2025). recalculate_all_logprobs: Whether to re-evaluate all logprobs. Returns: @@ -389,6 +403,7 @@ def get_scores( is_terminal_mask, sink_states_mask, i, + log_rewards=log_rewards, ) flattening_mask = trajectories.terminating_idx.lt( @@ -559,6 +574,7 @@ def loss( self, env: Env, trajectories: Trajectories, + log_rewards: torch.Tensor | None = None, recalculate_all_logprobs: bool = True, reduction: str = "mean", ) -> torch.Tensor: @@ -567,6 +583,11 @@ def loss( Args: env: The environment where the trajectories are sampled from. trajectories: The batch of trajectories to compute the loss with. + log_rewards: Optional custom log rewards tensor of shape + (n_trajectories,). When None, uses the environment rewards. + Useful for intrinsic rewards (see + "Towards Improving Exploration through Sibling Augmented + GFlowNets", Madan et al., ICLR 2025). recalculate_all_logprobs: Whether to re-evaluate all logprobs. reduction: The reduction method to use ('mean', 'sum', or 'none'). Note: for geometric-based sub-trajectory weighting, 'mean' is not @@ -598,12 +619,15 @@ def loss( logF_s0 = log_state_flows[0].masked_fill(log_state_flows[0].isinf(), 0.0) total_log_pf = log_pf_traj.sum(dim=0) total_log_pb = log_pb_traj.sum(dim=0) - # Guard behind debug: log_rewards is always set for terminating trajectories; - # bare assert would cause a graph break in torch.compile. + if log_rewards is None: + # Guard behind debug: log_rewards is always set for terminating + # trajectories; bare assert would cause a graph break in torch.compile. + if self.debug: + assert trajectories.log_rewards is not None + # cast: log_rewards is always set for terminating trajectories. + log_rewards = cast(torch.Tensor, trajectories.log_rewards) if self.debug: - assert trajectories.log_rewards is not None - # cast: log_rewards is always set for terminating trajectories. - log_rewards = cast(torch.Tensor, trajectories.log_rewards) + assert log_rewards.shape == (trajectories.batch_size,) if math.isfinite(self.log_reward_clip_min): log_rewards = log_rewards.clamp_min(self.log_reward_clip_min) tb_scores = logF_s0 + total_log_pf - total_log_pb - log_rewards @@ -611,7 +635,10 @@ def loss( # Get all scores and masks from the trajectories. scores, flattening_masks = self.get_scores( - trajectories, recalculate_all_logprobs=recalculate_all_logprobs, env=env + trajectories, + log_rewards=log_rewards, + recalculate_all_logprobs=recalculate_all_logprobs, + env=env, ) flattening_mask = torch.cat(flattening_masks) all_scores = torch.cat(scores, 0) diff --git a/src/gfn/gflownet/trajectory_balance.py b/src/gfn/gflownet/trajectory_balance.py index 1cea610b..58066622 100644 --- a/src/gfn/gflownet/trajectory_balance.py +++ b/src/gfn/gflownet/trajectory_balance.py @@ -91,6 +91,7 @@ def loss( self, env: Env, trajectories: Trajectories, + log_rewards: torch.Tensor | None = None, recalculate_all_logprobs: bool = True, reduction: str = "mean", ) -> torch.Tensor: @@ -102,6 +103,11 @@ def loss( Args: env: The environment where the trajectories are sampled from (unused). trajectories: The Trajectories object to compute the loss with. + log_rewards: Optional custom log rewards tensor of shape + (n_trajectories,). When None, uses the environment rewards. + Useful for intrinsic rewards (see + "Towards Improving Exploration through Sibling Augmented + GFlowNets", Madan et al., ICLR 2025). recalculate_all_logprobs: Whether to re-evaluate all logprobs. reduction: The reduction method to use ('mean', 'sum', or 'none'). @@ -112,7 +118,9 @@ def loss( if self.debug: warn_about_recalculating_logprobs(trajectories, recalculate_all_logprobs) scores = self.get_scores( - trajectories, recalculate_all_logprobs=recalculate_all_logprobs + trajectories, + log_rewards=log_rewards, + recalculate_all_logprobs=recalculate_all_logprobs, ) # If the conditions values exist, we pass them to self.logZ @@ -219,6 +227,7 @@ def prior_pf(self) -> Estimator: def get_scores( self, trajectories: Trajectories, + log_rewards: torch.Tensor | None = None, recalculate_all_logprobs: bool = True, env: Env | None = None, ) -> torch.Tensor: @@ -230,16 +239,29 @@ def get_scores( Returns: Shape ``(N,)`` per-trajectory scores. """ - return self._compute_rtb_scores(env, trajectories, recalculate_all_logprobs) + return self._compute_rtb_scores( + env, trajectories, log_rewards, recalculate_all_logprobs + ) def _compute_rtb_scores( self, env: Env | None, trajectories: Trajectories, + log_rewards: torch.Tensor | None = None, recalculate_all_logprobs: bool = True, ) -> torch.Tensor: """RTB residuals: ``log_pf_post - log_pf_prior - beta * log_rewards``. + Args: + env: The environment (unused, kept for API consistency). + trajectories: The Trajectories object to evaluate. + log_rewards: Optional custom log rewards tensor of shape + (n_trajectories,). When None, uses the environment rewards. + Useful for intrinsic rewards (see + "Towards Improving Exploration through Sibling Augmented + GFlowNets", Madan et al., ICLR 2025). + recalculate_all_logprobs: Whether to re-evaluate all logprobs. + Returns: Shape ``(N,)`` per-trajectory scores. """ @@ -262,9 +284,12 @@ def _compute_rtb_scores( recalculate_all_logprobs=True, ).sum(dim=0) + if log_rewards is None: + if self.debug: + assert trajectories.log_rewards is not None + log_rewards = cast(torch.Tensor, trajectories.log_rewards) if self.debug: - assert trajectories.log_rewards is not None - log_rewards = cast(torch.Tensor, trajectories.log_rewards) + assert log_rewards.shape == (trajectories.batch_size,) if math.isfinite(self.log_reward_clip_min): log_rewards = log_rewards.clamp_min(self.log_reward_clip_min) @@ -331,11 +356,17 @@ def loss( self, env: Env, trajectories: Trajectories, + log_rewards: torch.Tensor | None = None, recalculate_all_logprobs: bool = True, reduction: str = "mean", ) -> torch.Tensor: """Computes the RTB loss on a batch of trajectories.""" - scores = self._compute_rtb_scores(env, trajectories, recalculate_all_logprobs) + scores = self._compute_rtb_scores( + env, + trajectories, + log_rewards=log_rewards, + recalculate_all_logprobs=recalculate_all_logprobs, + ) # Get logZ. if trajectories.states.conditions is not None: @@ -547,6 +578,7 @@ def loss( self, env: Env, trajectories: Trajectories, + log_rewards: torch.Tensor | None = None, recalculate_all_logprobs: bool = True, reduction: str = "mean", ) -> torch.Tensor: @@ -556,7 +588,13 @@ def loss( :meth:`RelativeTrajectoryBalanceGFlowNet.loss`. It ensures gradient magnitudes match the Trust-PCL formulation. """ - rtb_loss = super().loss(env, trajectories, recalculate_all_logprobs, reduction) + rtb_loss = super().loss( + env, + trajectories, + log_rewards=log_rewards, + recalculate_all_logprobs=recalculate_all_logprobs, + reduction=reduction, + ) return self.alpha**2 * rtb_loss @@ -579,6 +617,7 @@ def loss( self, env: Env, trajectories: Trajectories, + log_rewards: torch.Tensor | None = None, recalculate_all_logprobs: bool = True, reduction: str = "mean", ) -> torch.Tensor: @@ -590,6 +629,11 @@ def loss( Args: env: The environment where the trajectories are sampled from (unused). trajectories: The Trajectories object to compute the loss with. + log_rewards: Optional custom log rewards tensor of shape + (n_trajectories,). When None, uses the environment rewards. + Useful for intrinsic rewards (see + "Towards Improving Exploration through Sibling Augmented + GFlowNets", Madan et al., ICLR 2025). recalculate_all_logprobs: Whether to re-evaluate all logprobs. reduction: The reduction method to use ('mean', 'sum', or 'none'). @@ -600,7 +644,9 @@ def loss( if self.debug: warn_about_recalculating_logprobs(trajectories, recalculate_all_logprobs) scores = self.get_scores( - trajectories, recalculate_all_logprobs=recalculate_all_logprobs + trajectories, + log_rewards=log_rewards, + recalculate_all_logprobs=recalculate_all_logprobs, ) scores = scores.sub_(scores.mean()) # [N], in-place mean-centering. scores = self.loss_fn(scores) # [N] @@ -663,11 +709,17 @@ def loss( self, env: Env, trajectories: Trajectories, + log_rewards: torch.Tensor | None = None, recalculate_all_logprobs: bool = True, reduction: str = "mean", ) -> torch.Tensor: """Computes the Relative LPV loss on a batch of trajectories.""" - scores = self._compute_rtb_scores(env, trajectories, recalculate_all_logprobs) + scores = self._compute_rtb_scores( + env, + trajectories, + log_rewards=log_rewards, + recalculate_all_logprobs=recalculate_all_logprobs, + ) scores = scores - scores.mean() scores = self.loss_fn(scores) diff --git a/tutorials/examples/train_hypergrid_gafn.py b/tutorials/examples/train_hypergrid_gafn.py index 46687bea..936cdc7b 100644 --- a/tutorials/examples/train_hypergrid_gafn.py +++ b/tutorials/examples/train_hypergrid_gafn.py @@ -239,10 +239,17 @@ def loss( self, env: Env, trajectories: Trajectories, + log_rewards: torch.Tensor | None = None, recalculate_all_logprobs: bool = True, reduction: str = "mean", ) -> torch.Tensor: - loss = super().loss(env, trajectories, recalculate_all_logprobs, reduction) + loss = super().loss( + env, + trajectories, + log_rewards=log_rewards, + recalculate_all_logprobs=recalculate_all_logprobs, + reduction=reduction, + ) rnd_loss = self.rnd.compute_rnd_loss(trajectories.states) return loss + rnd_loss From 9630594bd43b47e5008ba704da9d9c1d294971e0 Mon Sep 17 00:00:00 2001 From: Joseph Viviano Date: Thu, 2 Apr 2026 22:47:24 -0400 Subject: [PATCH 02/10] Add end-to-end tests for custom log_rewards parameter 14 tests covering all GFlowNet loss variants (TB, DB, SubTB, FM, LPV, RTB) verify two properties: 1. Identity: passing log_rewards equal to the environment rewards reproduces the default loss exactly. 2. Override: passing different log_rewards produces a different loss, proving the custom rewards are actually used. Co-Authored-By: Claude Opus 4.6 (1M context) --- testing/test_custom_log_rewards.py | 368 +++++++++++++++++++++++++++++ 1 file changed, 368 insertions(+) create mode 100644 testing/test_custom_log_rewards.py diff --git a/testing/test_custom_log_rewards.py b/testing/test_custom_log_rewards.py new file mode 100644 index 00000000..ab3ce63f --- /dev/null +++ b/testing/test_custom_log_rewards.py @@ -0,0 +1,368 @@ +"""Tests that the optional `log_rewards` parameter correctly overrides +environment rewards in all GFlowNet loss implementations. + +For each loss variant we verify two properties: + +1. **Identity**: passing `log_rewards=trajectories.log_rewards` (or the + equivalent for transition/state-based losses) reproduces the default loss + exactly. +2. **Override**: passing a *different* `log_rewards` tensor produces a + different loss value, proving that the custom rewards are actually used. +""" + +import torch + +from gfn.estimators import DiscretePolicyEstimator, ScalarEstimator +from gfn.gflownet.detailed_balance import DBGFlowNet +from gfn.gflownet.flow_matching import FMGFlowNet +from gfn.gflownet.sub_trajectory_balance import SubTBGFlowNet +from gfn.gflownet.trajectory_balance import ( + LogPartitionVarianceGFlowNet, + RelativeTrajectoryBalanceGFlowNet, + TBGFlowNet, +) +from gfn.gym import HyperGrid +from gfn.preprocessors import KHotPreprocessor +from gfn.samplers import Sampler +from gfn.utils.modules import MLP + +# --------------------------------------------------------------------------- +# Shared helpers +# --------------------------------------------------------------------------- + + +def _make_hypergrid_estimators(): + """Build a HyperGrid env with PF, PB, and logF estimators.""" + torch.manual_seed(0) + env = HyperGrid(ndim=2, height=4, validate_modes=False) + preproc = KHotPreprocessor(env.height, env.ndim) + assert isinstance(preproc.output_dim, int) + + pf_module = MLP(input_dim=preproc.output_dim, output_dim=env.n_actions) + pb_module = MLP(input_dim=preproc.output_dim, output_dim=env.n_actions - 1) + logF_module = MLP(input_dim=preproc.output_dim, output_dim=1) + + pf = DiscretePolicyEstimator( + module=pf_module, + n_actions=env.n_actions, + preprocessor=preproc, + is_backward=False, + ) + pb = DiscretePolicyEstimator( + module=pb_module, + n_actions=env.n_actions, + preprocessor=preproc, + is_backward=True, + ) + logF = ScalarEstimator(module=logF_module, preprocessor=preproc) + return env, pf, pb, logF + + +def _sample_trajectories(env, pf, n=8): + """Sample trajectories from the environment.""" + sampler = Sampler(estimator=pf) + return sampler.sample_trajectories(env, n=n, save_logprobs=True) + + +# --------------------------------------------------------------------------- +# TB +# --------------------------------------------------------------------------- + + +class TestTBCustomLogRewards: + def test_identity(self): + """Passing the trajectory's own log_rewards reproduces the default loss.""" + env, pf, pb, _ = _make_hypergrid_estimators() + gfn = TBGFlowNet(pf=pf, pb=pb) + trajs = _sample_trajectories(env, pf) + + with torch.no_grad(): + default_loss = gfn.loss(env, trajs) + explicit_loss = gfn.loss(env, trajs, log_rewards=trajs.log_rewards) + + assert torch.allclose(default_loss, explicit_loss), ( + f"Identity failed: default={default_loss.item()}, " + f"explicit={explicit_loss.item()}" + ) + + def test_override(self): + """Passing different log_rewards produces a different loss.""" + env, pf, pb, _ = _make_hypergrid_estimators() + gfn = TBGFlowNet(pf=pf, pb=pb) + trajs = _sample_trajectories(env, pf) + + custom_rewards = torch.zeros(trajs.batch_size) + + with torch.no_grad(): + default_loss = gfn.loss(env, trajs) + custom_loss = gfn.loss(env, trajs, log_rewards=custom_rewards) + + assert not torch.allclose( + default_loss, custom_loss + ), "Override failed: custom rewards produced the same loss as default" + + +# --------------------------------------------------------------------------- +# DB +# --------------------------------------------------------------------------- + + +class TestDBCustomLogRewards: + def test_identity(self): + """Passing the transition's own log_rewards reproduces the default loss.""" + env, pf, pb, logF = _make_hypergrid_estimators() + gfn = DBGFlowNet(pf=pf, pb=pb, logF=logF) + trajs = _sample_trajectories(env, pf) + transitions = trajs.to_transitions() + + with torch.no_grad(): + default_loss = gfn.loss(env, transitions) + explicit_loss = gfn.loss( + env, transitions, log_rewards=transitions.log_rewards + ) + + assert torch.allclose(default_loss, explicit_loss), ( + f"Identity failed: default={default_loss.item()}, " + f"explicit={explicit_loss.item()}" + ) + + def test_override(self): + """Passing different log_rewards produces a different loss.""" + env, pf, pb, logF = _make_hypergrid_estimators() + gfn = DBGFlowNet(pf=pf, pb=pb, logF=logF) + trajs = _sample_trajectories(env, pf) + transitions = trajs.to_transitions() + + custom_rewards = torch.zeros(transitions.n_transitions) + + with torch.no_grad(): + default_loss = gfn.loss(env, transitions) + custom_loss = gfn.loss(env, transitions, log_rewards=custom_rewards) + + assert not torch.allclose( + default_loss, custom_loss + ), "Override failed: custom rewards produced the same loss as default" + + +# --------------------------------------------------------------------------- +# SubTB +# --------------------------------------------------------------------------- + + +class TestSubTBCustomLogRewards: + def test_identity(self): + """Passing the trajectory's own log_rewards reproduces the default loss.""" + env, pf, pb, logF = _make_hypergrid_estimators() + gfn = SubTBGFlowNet(pf=pf, pb=pb, logF=logF, weighting="geometric_within") + trajs = _sample_trajectories(env, pf) + + with torch.no_grad(): + default_loss = gfn.loss(env, trajs) + explicit_loss = gfn.loss(env, trajs, log_rewards=trajs.log_rewards) + + assert torch.allclose(default_loss, explicit_loss), ( + f"Identity failed: default={default_loss.item()}, " + f"explicit={explicit_loss.item()}" + ) + + def test_override(self): + """Passing different log_rewards produces a different loss.""" + env, pf, pb, logF = _make_hypergrid_estimators() + gfn = SubTBGFlowNet(pf=pf, pb=pb, logF=logF, weighting="geometric_within") + trajs = _sample_trajectories(env, pf) + + custom_rewards = torch.zeros(trajs.batch_size) + + with torch.no_grad(): + default_loss = gfn.loss(env, trajs) + custom_loss = gfn.loss(env, trajs, log_rewards=custom_rewards) + + assert not torch.allclose( + default_loss, custom_loss + ), "Override failed: custom rewards produced the same loss as default" + + def test_identity_tb_weighting(self): + """TB weighting path also respects custom log_rewards identity.""" + env, pf, pb, logF = _make_hypergrid_estimators() + gfn = SubTBGFlowNet(pf=pf, pb=pb, logF=logF, weighting="TB") + trajs = _sample_trajectories(env, pf) + + with torch.no_grad(): + default_loss = gfn.loss(env, trajs) + explicit_loss = gfn.loss(env, trajs, log_rewards=trajs.log_rewards) + + assert torch.allclose(default_loss, explicit_loss), ( + f"Identity (TB weighting) failed: default={default_loss.item()}, " + f"explicit={explicit_loss.item()}" + ) + + def test_override_tb_weighting(self): + """TB weighting path also uses custom log_rewards when provided.""" + env, pf, pb, logF = _make_hypergrid_estimators() + gfn = SubTBGFlowNet(pf=pf, pb=pb, logF=logF, weighting="TB") + trajs = _sample_trajectories(env, pf) + + custom_rewards = torch.zeros(trajs.batch_size) + + with torch.no_grad(): + default_loss = gfn.loss(env, trajs) + custom_loss = gfn.loss(env, trajs, log_rewards=custom_rewards) + + assert not torch.allclose( + default_loss, custom_loss + ), "Override (TB weighting) failed: custom rewards produced the same loss" + + +# --------------------------------------------------------------------------- +# FM +# --------------------------------------------------------------------------- + + +class TestFMCustomLogRewards: + def test_identity(self): + """Passing the container's own terminating_log_rewards reproduces the + default loss.""" + torch.manual_seed(0) + env = HyperGrid(ndim=2, height=4, validate_modes=False) + preproc = KHotPreprocessor(env.height, env.ndim) + module = MLP( + input_dim=preproc.output_dim, output_dim=env.n_actions # type: ignore + ) + logF = DiscretePolicyEstimator( + module=module, n_actions=env.n_actions, preprocessor=preproc + ) + gfn = FMGFlowNet(logF) + + trajs = gfn.sample_trajectories(env, n=8) + states_container = gfn.to_training_samples(trajs) + + with torch.no_grad(): + default_loss = gfn.loss(env, states_container) + explicit_loss = gfn.loss( + env, + states_container, + log_rewards=states_container.terminating_log_rewards, + ) + + assert torch.allclose(default_loss, explicit_loss), ( + f"Identity failed: default={default_loss.item()}, " + f"explicit={explicit_loss.item()}" + ) + + def test_override(self): + """Passing different log_rewards produces a different loss.""" + torch.manual_seed(0) + env = HyperGrid(ndim=2, height=4, validate_modes=False) + preproc = KHotPreprocessor(env.height, env.ndim) + module = MLP( + input_dim=preproc.output_dim, output_dim=env.n_actions # type: ignore + ) + logF = DiscretePolicyEstimator( + module=module, n_actions=env.n_actions, preprocessor=preproc + ) + gfn = FMGFlowNet(logF) + + trajs = gfn.sample_trajectories(env, n=8) + states_container = gfn.to_training_samples(trajs) + + n_terminating = len(states_container.terminating_states) + custom_rewards = torch.zeros(n_terminating) + + with torch.no_grad(): + default_loss = gfn.loss(env, states_container) + custom_loss = gfn.loss(env, states_container, log_rewards=custom_rewards) + + assert not torch.allclose( + default_loss, custom_loss + ), "Override failed: custom rewards produced the same loss as default" + + +# --------------------------------------------------------------------------- +# LPV +# --------------------------------------------------------------------------- + + +class TestLPVCustomLogRewards: + def test_identity(self): + """Passing the trajectory's own log_rewards reproduces the default loss.""" + env, pf, pb, _ = _make_hypergrid_estimators() + gfn = LogPartitionVarianceGFlowNet(pf=pf, pb=pb) + trajs = _sample_trajectories(env, pf) + + with torch.no_grad(): + default_loss = gfn.loss(env, trajs) + explicit_loss = gfn.loss(env, trajs, log_rewards=trajs.log_rewards) + + assert torch.allclose(default_loss, explicit_loss), ( + f"Identity failed: default={default_loss.item()}, " + f"explicit={explicit_loss.item()}" + ) + + def test_override(self): + """Passing different log_rewards produces a different loss.""" + env, pf, pb, _ = _make_hypergrid_estimators() + gfn = LogPartitionVarianceGFlowNet(pf=pf, pb=pb) + trajs = _sample_trajectories(env, pf) + + custom_rewards = torch.zeros(trajs.batch_size) + + with torch.no_grad(): + default_loss = gfn.loss(env, trajs) + custom_loss = gfn.loss(env, trajs, log_rewards=custom_rewards) + + assert not torch.allclose( + default_loss, custom_loss + ), "Override failed: custom rewards produced the same loss as default" + + +# --------------------------------------------------------------------------- +# RTB +# --------------------------------------------------------------------------- + + +class TestRTBCustomLogRewards: + def _make_rtb(self): + env, pf, _, _ = _make_hypergrid_estimators() + preproc = KHotPreprocessor(env.height, env.ndim) + prior_module = MLP( + input_dim=preproc.output_dim, output_dim=env.n_actions # type: ignore + ) + prior_pf = DiscretePolicyEstimator( + module=prior_module, + n_actions=env.n_actions, + preprocessor=preproc, + is_backward=False, + ) + for p in prior_pf.parameters(): + p.requires_grad_(False) + + gfn = RelativeTrajectoryBalanceGFlowNet(pf=pf, prior_pf=prior_pf, beta=1.0) + trajs = _sample_trajectories(env, pf) + return env, gfn, trajs + + def test_identity(self): + """Passing the trajectory's own log_rewards reproduces the default loss.""" + env, gfn, trajs = self._make_rtb() + + with torch.no_grad(): + default_loss = gfn.loss(env, trajs) + explicit_loss = gfn.loss(env, trajs, log_rewards=trajs.log_rewards) + + assert torch.allclose(default_loss, explicit_loss), ( + f"Identity failed: default={default_loss.item()}, " + f"explicit={explicit_loss.item()}" + ) + + def test_override(self): + """Passing different log_rewards produces a different loss.""" + env, gfn, trajs = self._make_rtb() + + custom_rewards = torch.zeros(trajs.batch_size) + + with torch.no_grad(): + default_loss = gfn.loss(env, trajs) + custom_loss = gfn.loss(env, trajs, log_rewards=custom_rewards) + + assert not torch.allclose( + default_loss, custom_loss + ), "Override failed: custom rewards produced the same loss as default" From 2308c246d33f9c8478bc9aefc1bdc3e23ca84db5 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 3 Apr 2026 04:25:11 +0000 Subject: [PATCH 03/10] Initial plan From 485dab72173b5db10a89c30c9c649378e092035c Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 3 Apr 2026 04:25:51 +0000 Subject: [PATCH 04/10] fix: accept log_rewards kwarg in TBGAFN.get_scores with default None Agent-Logs-Url: https://github.com/GFNOrg/torchgfn/sessions/1919552f-66e6-465f-9209-ebd7b2f8629a Co-authored-by: josephdviviano <4142570+josephdviviano@users.noreply.github.com> --- tutorials/examples/train_hypergrid_gafn.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tutorials/examples/train_hypergrid_gafn.py b/tutorials/examples/train_hypergrid_gafn.py index 936cdc7b..7ad43cc6 100644 --- a/tutorials/examples/train_hypergrid_gafn.py +++ b/tutorials/examples/train_hypergrid_gafn.py @@ -162,7 +162,10 @@ def flow_parameters(self) -> list[torch.Tensor]: ) def get_scores( - self, trajectories: Trajectories, recalculate_all_logprobs: bool = True + self, + trajectories: Trajectories, + recalculate_all_logprobs: bool = True, + log_rewards: torch.Tensor | None = None, ) -> torch.Tensor: """Computes Trajectory Balance scores with intrinsic rewards for a batch of trajectories. @@ -170,6 +173,8 @@ def get_scores( Args: trajectories: The Trajectories object to evaluate. recalculate_all_logprobs: Whether to re-evaluate all logprobs. + log_rewards: Optional override for the trajectories' log rewards. If None, + defaults to `trajectories.log_rewards`. Returns: A tensor of shape (batch_size,) containing the scores for each trajectory. @@ -177,7 +182,8 @@ def get_scores( log_pf_trajectories, log_pb_trajectories = self.get_pfs_and_pbs( trajectories, recalculate_all_logprobs=recalculate_all_logprobs ) - log_rewards = trajectories.log_rewards + if log_rewards is None: + log_rewards = trajectories.log_rewards assert log_rewards is not None if self.use_edge_ri: # Use the edge-based intrinsic rewards. From ee6f7872c9026a6875736eb9e56e609081972108 Mon Sep 17 00:00:00 2001 From: Joseph Viviano Date: Fri, 3 Apr 2026 01:54:07 -0400 Subject: [PATCH 05/10] Update src/gfn/gflownet/sub_trajectory_balance.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- src/gfn/gflownet/sub_trajectory_balance.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/gfn/gflownet/sub_trajectory_balance.py b/src/gfn/gflownet/sub_trajectory_balance.py index eab226ab..f4e4a9b4 100644 --- a/src/gfn/gflownet/sub_trajectory_balance.py +++ b/src/gfn/gflownet/sub_trajectory_balance.py @@ -585,9 +585,13 @@ def loss( trajectories: The batch of trajectories to compute the loss with. log_rewards: Optional custom log rewards tensor of shape (n_trajectories,). When None, uses the environment rewards. - Useful for intrinsic rewards (see - "Towards Improving Exploration through Sibling Augmented - GFlowNets", Madan et al., ICLR 2025). + When provided, this overrides the terminal reward term used by + the loss. In particular, for ``forward_looking=True``, the + state-flow computation may still depend on ``env.log_reward(...)``, + so custom ``log_rewards`` do not fully replace environment + rewards in that mode. Useful for intrinsic rewards affecting the + terminal boundary term (see "Towards Improving Exploration + through Sibling Augmented GFlowNets", Madan et al., ICLR 2025). recalculate_all_logprobs: Whether to re-evaluate all logprobs. reduction: The reduction method to use ('mean', 'sum', or 'none'). Note: for geometric-based sub-trajectory weighting, 'mean' is not From 11caee99e2c1bdf0617601a2aef52fa55caa40a7 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 3 Apr 2026 06:04:18 +0000 Subject: [PATCH 06/10] fix: make log_rewards keyword-only in all loss() and get_scores() methods for backward compatibility Agent-Logs-Url: https://github.com/GFNOrg/torchgfn/sessions/59f54f6c-8302-486a-ae28-3b37747a3b55 Co-authored-by: josephdviviano <4142570+josephdviviano@users.noreply.github.com> --- src/gfn/gflownet/base.py | 9 +++---- src/gfn/gflownet/detailed_balance.py | 12 ++++++---- src/gfn/gflownet/flow_matching.py | 7 +++--- src/gfn/gflownet/sub_trajectory_balance.py | 20 +++++++++------- src/gfn/gflownet/trajectory_balance.py | 28 +++++++++++++--------- tutorials/examples/train_hypergrid_gafn.py | 4 +++- 6 files changed, 47 insertions(+), 33 deletions(-) diff --git a/src/gfn/gflownet/base.py b/src/gfn/gflownet/base.py index 3ac0d34f..c152677e 100644 --- a/src/gfn/gflownet/base.py +++ b/src/gfn/gflownet/base.py @@ -420,9 +420,10 @@ def logz_parameters(self) -> list[torch.Tensor]: def get_scores( self, trajectories: Trajectories, - log_rewards: torch.Tensor | None = None, recalculate_all_logprobs: bool = True, env: Env | None = None, + *, + log_rewards: torch.Tensor | None = None, ) -> torch.Tensor: r"""Calculates scores for a batch of trajectories. @@ -431,14 +432,14 @@ def get_scores( Args: trajectories: The Trajectories object to evaluate. + recalculate_all_logprobs: Whether to re-evaluate all logprobs. + env: The environment (unused in base TB, but required by some + subclasses such as RTB and SubTB). log_rewards: Optional custom log rewards tensor of shape (n_trajectories,). When None, uses the environment rewards from the trajectories. Useful for intrinsic rewards (see "Towards Improving Exploration through Sibling Augmented GFlowNets", Madan et al., ICLR 2025). - recalculate_all_logprobs: Whether to re-evaluate all logprobs. - env: The environment (unused in base TB, but required by some - subclasses such as RTB and SubTB). Returns: A tensor of shape (batch_size,) containing the scores for each trajectory. diff --git a/src/gfn/gflownet/detailed_balance.py b/src/gfn/gflownet/detailed_balance.py index 2b43638e..5ebaa020 100644 --- a/src/gfn/gflownet/detailed_balance.py +++ b/src/gfn/gflownet/detailed_balance.py @@ -209,8 +209,9 @@ def get_scores( self, env: Env, transitions: Transitions, - log_rewards: torch.Tensor | None = None, recalculate_all_logprobs: bool = True, + *, + log_rewards: torch.Tensor | None = None, ) -> torch.Tensor: r"""Calculates the scores for a batch of transitions. @@ -220,12 +221,12 @@ def get_scores( Args: env: The environment where the transitions are sampled from. transitions: The Transitions object to evaluate. + recalculate_all_logprobs: Whether to re-evaluate all logprobs. log_rewards: Optional custom log rewards tensor of shape (n_transitions,). When None, uses the environment rewards from the transitions. Useful for intrinsic rewards (see "Towards Improving Exploration through Sibling Augmented GFlowNets", Madan et al., ICLR 2025). - recalculate_all_logprobs: Whether to re-evaluate all logprobs. Returns: A tensor of shape (n_transitions,) representing the scores for each @@ -340,9 +341,10 @@ def loss( self, env: Env, transitions: Transitions, - log_rewards: torch.Tensor | None = None, recalculate_all_logprobs: bool = True, reduction: str = "mean", + *, + log_rewards: torch.Tensor | None = None, ) -> torch.Tensor: """Computes the detailed balance loss. @@ -352,11 +354,11 @@ def loss( Args: env: The environment where the transitions are sampled from. transitions: The Transitions object to compute the loss with. - log_rewards: Optional custom log rewards tensor of shape - (n_transitions,). When None, uses the environment rewards. recalculate_all_logprobs: Whether to re-evaluate all logprobs. reduction: The reduction method to use ('mean', 'sum', or 'none'). Run with self.debug=False for improved performance. + log_rewards: Optional custom log rewards tensor of shape + (n_transitions,). When None, uses the environment rewards. Returns: The computed detailed balance loss as a tensor. The shape depends on the diff --git a/src/gfn/gflownet/flow_matching.py b/src/gfn/gflownet/flow_matching.py index 24edec27..713df950 100644 --- a/src/gfn/gflownet/flow_matching.py +++ b/src/gfn/gflownet/flow_matching.py @@ -279,9 +279,10 @@ def loss( self, env: DiscreteEnv, states_container: StatesContainer[DiscreteStates], - log_rewards: torch.Tensor | None = None, recalculate_all_logprobs: bool = True, reduction: str = "mean", + *, + log_rewards: torch.Tensor | None = None, ) -> torch.Tensor: """Computes the flow matching loss for a batch of states. @@ -294,13 +295,13 @@ def loss( env: The discrete environment where the states are sampled from. states_container: The StatesContainer object containing both intermediary and terminating states. + recalculate_all_logprobs: Whether to re-evaluate all logprobs (unused for FM). + reduction: The reduction method to use ('mean', 'sum', or 'none'). log_rewards: Optional custom log rewards tensor of shape (n_terminating_states,). When None, uses the environment rewards from the states container. Useful for intrinsic rewards (see "Towards Improving Exploration through Sibling Augmented GFlowNets", Madan et al., ICLR 2025). - recalculate_all_logprobs: Whether to re-evaluate all logprobs (unused for FM). - reduction: The reduction method to use ('mean', 'sum', or 'none'). Returns: The computed flow matching loss as a tensor. The shape depends on the diff --git a/src/gfn/gflownet/sub_trajectory_balance.py b/src/gfn/gflownet/sub_trajectory_balance.py index f4e4a9b4..0b7ebf53 100644 --- a/src/gfn/gflownet/sub_trajectory_balance.py +++ b/src/gfn/gflownet/sub_trajectory_balance.py @@ -341,21 +341,22 @@ def calculate_masks( def get_scores( self, trajectories: Trajectories, - log_rewards: torch.Tensor | None = None, recalculate_all_logprobs: bool = True, env: Env | None = None, + *, + log_rewards: torch.Tensor | None = None, ) -> Tuple[List[torch.Tensor], List[torch.Tensor]]: r"""Computes sub-trajectory balance scores for all submitted trajectories. Args: - env: The environment where the trajectories are sampled from. trajectories: The batch of trajectories to evaluate. + recalculate_all_logprobs: Whether to re-evaluate all logprobs. + env: The environment where the trajectories are sampled from. log_rewards: Optional custom log rewards tensor of shape (n_trajectories,). When None, uses the environment rewards. Useful for intrinsic rewards (see "Towards Improving Exploration through Sibling Augmented GFlowNets", Madan et al., ICLR 2025). - recalculate_all_logprobs: Whether to re-evaluate all logprobs. Returns: A tuple (scores, flattening_masks): @@ -574,15 +575,21 @@ def loss( self, env: Env, trajectories: Trajectories, - log_rewards: torch.Tensor | None = None, recalculate_all_logprobs: bool = True, reduction: str = "mean", + *, + log_rewards: torch.Tensor | None = None, ) -> torch.Tensor: """Computes the sub-trajectory balance loss. Args: env: The environment where the trajectories are sampled from. trajectories: The batch of trajectories to compute the loss with. + recalculate_all_logprobs: Whether to re-evaluate all logprobs. + reduction: The reduction method to use ('mean', 'sum', or 'none'). + Note: for geometric-based sub-trajectory weighting, 'mean' is not + supported and is coerced to 'sum' (a warning is emitted when + debug=True). log_rewards: Optional custom log rewards tensor of shape (n_trajectories,). When None, uses the environment rewards. When provided, this overrides the terminal reward term used by @@ -592,11 +599,6 @@ def loss( rewards in that mode. Useful for intrinsic rewards affecting the terminal boundary term (see "Towards Improving Exploration through Sibling Augmented GFlowNets", Madan et al., ICLR 2025). - recalculate_all_logprobs: Whether to re-evaluate all logprobs. - reduction: The reduction method to use ('mean', 'sum', or 'none'). - Note: for geometric-based sub-trajectory weighting, 'mean' is not - supported and is coerced to 'sum' (a warning is emitted when - debug=True). Returns: The computed sub-trajectory balance loss as a tensor. The shape depends on diff --git a/src/gfn/gflownet/trajectory_balance.py b/src/gfn/gflownet/trajectory_balance.py index 58066622..b3405245 100644 --- a/src/gfn/gflownet/trajectory_balance.py +++ b/src/gfn/gflownet/trajectory_balance.py @@ -91,9 +91,10 @@ def loss( self, env: Env, trajectories: Trajectories, - log_rewards: torch.Tensor | None = None, recalculate_all_logprobs: bool = True, reduction: str = "mean", + *, + log_rewards: torch.Tensor | None = None, ) -> torch.Tensor: """Computes the trajectory balance loss. @@ -103,13 +104,13 @@ def loss( Args: env: The environment where the trajectories are sampled from (unused). trajectories: The Trajectories object to compute the loss with. + recalculate_all_logprobs: Whether to re-evaluate all logprobs. + reduction: The reduction method to use ('mean', 'sum', or 'none'). log_rewards: Optional custom log rewards tensor of shape (n_trajectories,). When None, uses the environment rewards. Useful for intrinsic rewards (see "Towards Improving Exploration through Sibling Augmented GFlowNets", Madan et al., ICLR 2025). - recalculate_all_logprobs: Whether to re-evaluate all logprobs. - reduction: The reduction method to use ('mean', 'sum', or 'none'). Returns: The computed trajectory balance loss as a tensor. The shape depends on the @@ -227,9 +228,10 @@ def prior_pf(self) -> Estimator: def get_scores( self, trajectories: Trajectories, - log_rewards: torch.Tensor | None = None, recalculate_all_logprobs: bool = True, env: Env | None = None, + *, + log_rewards: torch.Tensor | None = None, ) -> torch.Tensor: """RTB residuals (without logZ): ``log_pf_post - log_pf_prior - beta * log_R``. @@ -240,7 +242,7 @@ def get_scores( Shape ``(N,)`` per-trajectory scores. """ return self._compute_rtb_scores( - env, trajectories, log_rewards, recalculate_all_logprobs + env, trajectories, log_rewards=log_rewards, recalculate_all_logprobs=recalculate_all_logprobs ) def _compute_rtb_scores( @@ -356,9 +358,10 @@ def loss( self, env: Env, trajectories: Trajectories, - log_rewards: torch.Tensor | None = None, recalculate_all_logprobs: bool = True, reduction: str = "mean", + *, + log_rewards: torch.Tensor | None = None, ) -> torch.Tensor: """Computes the RTB loss on a batch of trajectories.""" scores = self._compute_rtb_scores( @@ -578,9 +581,10 @@ def loss( self, env: Env, trajectories: Trajectories, - log_rewards: torch.Tensor | None = None, recalculate_all_logprobs: bool = True, reduction: str = "mean", + *, + log_rewards: torch.Tensor | None = None, ) -> torch.Tensor: r"""Computes the Trust-PCL loss: :math:`\alpha^2 \cdot \mathcal{L}_{\text{RTB}}`. @@ -617,9 +621,10 @@ def loss( self, env: Env, trajectories: Trajectories, - log_rewards: torch.Tensor | None = None, recalculate_all_logprobs: bool = True, reduction: str = "mean", + *, + log_rewards: torch.Tensor | None = None, ) -> torch.Tensor: """Computes the log partition variance loss. @@ -629,13 +634,13 @@ def loss( Args: env: The environment where the trajectories are sampled from (unused). trajectories: The Trajectories object to compute the loss with. + recalculate_all_logprobs: Whether to re-evaluate all logprobs. + reduction: The reduction method to use ('mean', 'sum', or 'none'). log_rewards: Optional custom log rewards tensor of shape (n_trajectories,). When None, uses the environment rewards. Useful for intrinsic rewards (see "Towards Improving Exploration through Sibling Augmented GFlowNets", Madan et al., ICLR 2025). - recalculate_all_logprobs: Whether to re-evaluate all logprobs. - reduction: The reduction method to use ('mean', 'sum', or 'none'). Returns: The computed log partition variance loss as a tensor. The shape depends on @@ -709,9 +714,10 @@ def loss( self, env: Env, trajectories: Trajectories, - log_rewards: torch.Tensor | None = None, recalculate_all_logprobs: bool = True, reduction: str = "mean", + *, + log_rewards: torch.Tensor | None = None, ) -> torch.Tensor: """Computes the Relative LPV loss on a batch of trajectories.""" scores = self._compute_rtb_scores( diff --git a/tutorials/examples/train_hypergrid_gafn.py b/tutorials/examples/train_hypergrid_gafn.py index 7ad43cc6..76398e74 100644 --- a/tutorials/examples/train_hypergrid_gafn.py +++ b/tutorials/examples/train_hypergrid_gafn.py @@ -165,6 +165,7 @@ def get_scores( self, trajectories: Trajectories, recalculate_all_logprobs: bool = True, + *, log_rewards: torch.Tensor | None = None, ) -> torch.Tensor: """Computes Trajectory Balance scores with intrinsic rewards for a batch of @@ -245,9 +246,10 @@ def loss( self, env: Env, trajectories: Trajectories, - log_rewards: torch.Tensor | None = None, recalculate_all_logprobs: bool = True, reduction: str = "mean", + *, + log_rewards: torch.Tensor | None = None, ) -> torch.Tensor: loss = super().loss( env, From f7f0d408f24f4cdf59aab8fff1c2849a259576ae Mon Sep 17 00:00:00 2001 From: Joseph Viviano Date: Fri, 3 Apr 2026 02:31:31 -0400 Subject: [PATCH 07/10] Update testing/test_custom_log_rewards.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- testing/test_custom_log_rewards.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/testing/test_custom_log_rewards.py b/testing/test_custom_log_rewards.py index ab3ce63f..43832c94 100644 --- a/testing/test_custom_log_rewards.py +++ b/testing/test_custom_log_rewards.py @@ -1,5 +1,6 @@ """Tests that the optional `log_rewards` parameter correctly overrides -environment rewards in all GFlowNet loss implementations. +environment rewards in the TB/DB/SubTB/FM/LPV/RTB GFlowNet loss implementations +covered by this file. For each loss variant we verify two properties: From 001ecd39f8ababf3684425213af5e62e30d24a2a Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 3 Apr 2026 06:42:27 +0000 Subject: [PATCH 08/10] feat: raise ValueError when forward_looking=True with custom log_rewards; add TrustPCL and RelativeLPV tests Agent-Logs-Url: https://github.com/GFNOrg/torchgfn/sessions/88122379-1c44-4cfc-b15a-8b72a521b080 Co-authored-by: josephdviviano <4142570+josephdviviano@users.noreply.github.com> --- src/gfn/gflownet/detailed_balance.py | 15 +++ testing/test_custom_log_rewards.py | 159 ++++++++++++++++++++++++++- 2 files changed, 172 insertions(+), 2 deletions(-) diff --git a/src/gfn/gflownet/detailed_balance.py b/src/gfn/gflownet/detailed_balance.py index 5ebaa020..55c4e6d9 100644 --- a/src/gfn/gflownet/detailed_balance.py +++ b/src/gfn/gflownet/detailed_balance.py @@ -227,11 +227,26 @@ def get_scores( from the transitions. Useful for intrinsic rewards (see "Towards Improving Exploration through Sibling Augmented GFlowNets", Madan et al., ICLR 2025). + **Not supported when** ``forward_looking=True``: raises + ``ValueError`` in that case because the forward-looking + objective still calls ``env.log_reward()`` for intermediate + state adjustments, so custom rewards cannot fully replace + environment rewards. Returns: A tensor of shape (n_transitions,) representing the scores for each transition. """ + # Reject custom log_rewards in forward-looking mode (documented in docstring). + if log_rewards is not None and self.forward_looking: + raise ValueError( + "custom log_rewards are not supported when forward_looking=True: " + "the forward-looking DB objective calls env.log_reward() for " + "intermediate-state adjustments, so environment rewards cannot be " + "fully replaced. Either use forward_looking=False or do not pass " + "log_rewards." + ) + # Guard bad inputs under debug to avoid graph breaks in torch.compile. if self.debug and transitions.is_backward: raise ValueError("Backward transitions are not supported") diff --git a/testing/test_custom_log_rewards.py b/testing/test_custom_log_rewards.py index 43832c94..df4d2718 100644 --- a/testing/test_custom_log_rewards.py +++ b/testing/test_custom_log_rewards.py @@ -1,6 +1,6 @@ """Tests that the optional `log_rewards` parameter correctly overrides -environment rewards in the TB/DB/SubTB/FM/LPV/RTB GFlowNet loss implementations -covered by this file. +environment rewards in the TB/DB/SubTB/FM/RTB/LPV/TrustPCL/RelativeLPV +GFlowNet loss implementations covered by this file. For each loss variant we verify two properties: @@ -11,6 +11,7 @@ different loss value, proving that the custom rewards are actually used. """ +import pytest import torch from gfn.estimators import DiscretePolicyEstimator, ScalarEstimator @@ -19,8 +20,10 @@ from gfn.gflownet.sub_trajectory_balance import SubTBGFlowNet from gfn.gflownet.trajectory_balance import ( LogPartitionVarianceGFlowNet, + RelativeLogPartitionVarianceGFlowNet, RelativeTrajectoryBalanceGFlowNet, TBGFlowNet, + TrustPCLGFlowNet, ) from gfn.gym import HyperGrid from gfn.preprocessors import KHotPreprocessor @@ -144,6 +147,18 @@ def test_override(self): default_loss, custom_loss ), "Override failed: custom rewards produced the same loss as default" + def test_forward_looking_raises_with_custom_log_rewards(self): + """Providing log_rewards with forward_looking=True should raise ValueError.""" + env, pf, pb, logF = _make_hypergrid_estimators() + gfn = DBGFlowNet(pf=pf, pb=pb, logF=logF, forward_looking=True) + trajs = _sample_trajectories(env, pf) + transitions = trajs.to_transitions() + + custom_rewards = torch.zeros(transitions.n_transitions) + + with pytest.raises(ValueError, match="forward_looking"): + gfn.get_scores(env, transitions, log_rewards=custom_rewards) + # --------------------------------------------------------------------------- # SubTB @@ -367,3 +382,143 @@ def test_override(self): assert not torch.allclose( default_loss, custom_loss ), "Override failed: custom rewards produced the same loss as default" + + +# --------------------------------------------------------------------------- +# TrustPCL +# --------------------------------------------------------------------------- + + +class TestTrustPCLCustomLogRewards: + def _make_trust_pcl(self): + env, pf, _, _ = _make_hypergrid_estimators() + preproc = KHotPreprocessor(env.height, env.ndim) + prior_module = MLP( + input_dim=preproc.output_dim, output_dim=env.n_actions # type: ignore + ) + prior_pf = DiscretePolicyEstimator( + module=prior_module, + n_actions=env.n_actions, + preprocessor=preproc, + is_backward=False, + ) + for p in prior_pf.parameters(): + p.requires_grad_(False) + + gfn = TrustPCLGFlowNet(policy=pf, reference_policy=prior_pf, alpha=1.0) + trajs = _sample_trajectories(env, pf) + return env, gfn, trajs + + def test_identity(self): + """Passing the trajectory's own log_rewards reproduces the default loss.""" + env, gfn, trajs = self._make_trust_pcl() + + with torch.no_grad(): + default_loss = gfn.loss(env, trajs) + explicit_loss = gfn.loss(env, trajs, log_rewards=trajs.log_rewards) + + assert torch.allclose(default_loss, explicit_loss), ( + f"Identity failed: default={default_loss.item()}, " + f"explicit={explicit_loss.item()}" + ) + + def test_override(self): + """Passing different log_rewards produces a different loss.""" + env, gfn, trajs = self._make_trust_pcl() + + custom_rewards = torch.zeros(trajs.batch_size) + + with torch.no_grad(): + default_loss = gfn.loss(env, trajs) + custom_loss = gfn.loss(env, trajs, log_rewards=custom_rewards) + + assert not torch.allclose( + default_loss, custom_loss + ), "Override (TrustPCL) failed: custom rewards produced the same loss as default" + + def test_alpha_squared_scaling(self): + """TrustPCL loss is alpha^2 * RTB loss; custom rewards must survive scaling.""" + env, pf, _, _ = _make_hypergrid_estimators() + preproc = KHotPreprocessor(env.height, env.ndim) + prior_module = MLP( + input_dim=preproc.output_dim, output_dim=env.n_actions # type: ignore + ) + prior_pf = DiscretePolicyEstimator( + module=prior_module, + n_actions=env.n_actions, + preprocessor=preproc, + is_backward=False, + ) + for p in prior_pf.parameters(): + p.requires_grad_(False) + + alpha = 2.0 + gfn_rtb = RelativeTrajectoryBalanceGFlowNet( + pf=pf, prior_pf=prior_pf, beta=1.0 / alpha + ) + gfn_pcl = TrustPCLGFlowNet(policy=pf, reference_policy=prior_pf, alpha=alpha) + trajs = _sample_trajectories(env, pf) + + custom_rewards = torch.full((trajs.batch_size,), -1.0) + + with torch.no_grad(): + rtb_loss = gfn_rtb.loss(env, trajs, log_rewards=custom_rewards) + pcl_loss = gfn_pcl.loss(env, trajs, log_rewards=custom_rewards) + + assert torch.allclose(pcl_loss, alpha**2 * rtb_loss), ( + f"alpha^2 scaling failed: pcl={pcl_loss.item()}, " + f"alpha^2*rtb={alpha**2 * rtb_loss.item()}" + ) + + +# --------------------------------------------------------------------------- +# RelativeLPV +# --------------------------------------------------------------------------- + + +class TestRelativeLPVCustomLogRewards: + def _make_relative_lpv(self): + env, pf, _, _ = _make_hypergrid_estimators() + preproc = KHotPreprocessor(env.height, env.ndim) + prior_module = MLP( + input_dim=preproc.output_dim, output_dim=env.n_actions # type: ignore + ) + prior_pf = DiscretePolicyEstimator( + module=prior_module, + n_actions=env.n_actions, + preprocessor=preproc, + is_backward=False, + ) + for p in prior_pf.parameters(): + p.requires_grad_(False) + + gfn = RelativeLogPartitionVarianceGFlowNet(pf=pf, prior_pf=prior_pf, beta=1.0) + trajs = _sample_trajectories(env, pf) + return env, gfn, trajs + + def test_identity(self): + """Passing the trajectory's own log_rewards reproduces the default loss.""" + env, gfn, trajs = self._make_relative_lpv() + + with torch.no_grad(): + default_loss = gfn.loss(env, trajs) + explicit_loss = gfn.loss(env, trajs, log_rewards=trajs.log_rewards) + + assert torch.allclose(default_loss, explicit_loss), ( + f"Identity failed: default={default_loss.item()}, " + f"explicit={explicit_loss.item()}" + ) + + def test_override(self): + """Passing different log_rewards produces a different loss.""" + env, gfn, trajs = self._make_relative_lpv() + + custom_rewards = torch.zeros(trajs.batch_size) + + with torch.no_grad(): + default_loss = gfn.loss(env, trajs) + custom_loss = gfn.loss(env, trajs, log_rewards=custom_rewards) + + assert not torch.allclose( + default_loss, custom_loss + ), "Override (RelativeLPV) failed: custom rewards produced the same loss as default" From 380f59b2f6d8c2c3ece203cb78c2837cda7553e8 Mon Sep 17 00:00:00 2001 From: Joseph Viviano Date: Fri, 3 Apr 2026 03:30:34 -0400 Subject: [PATCH 09/10] Fix black formatting for CI (unstable string concatenation rules) Co-Authored-By: Claude Opus 4.6 (1M context) --- src/gfn/gflownet/trajectory_balance.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/gfn/gflownet/trajectory_balance.py b/src/gfn/gflownet/trajectory_balance.py index b3405245..ee3e2e4c 100644 --- a/src/gfn/gflownet/trajectory_balance.py +++ b/src/gfn/gflownet/trajectory_balance.py @@ -195,8 +195,8 @@ def __init__( raise ValueError( f"prior_pf has {len(trainable)} trainable parameter(s) " f"(first: {trainable[0]!r}). Freeze all prior parameters with " - f"`for p in prior_pf.parameters(): p.requires_grad_(False)` " - f"before constructing the RTB objective." + "`for p in prior_pf.parameters(): p.requires_grad_(False)` " + "before constructing the RTB objective." ) # Store the prior as a plain attribute (not an nn.Module submodule) @@ -511,7 +511,8 @@ def __init__( if not isinstance(logZ, type(None)) and init_logZ != 0.0: # If logZ is explicitly provided, we ignore init_v_soft_s0 and use logZ directly. warning.warn( - "TrustPCLGFlowNet's init_v_soft_s0 is ignored because logZ is explicitly provided. Ensure this is intentional." + "TrustPCLGFlowNet's init_v_soft_s0 is ignored because logZ is explicitly" + " provided. Ensure this is intentional." ) super().__init__( From 5aaef2e982391edd3a0fcb2db15e6ab49eda9324 Mon Sep 17 00:00:00 2001 From: Joseph Viviano Date: Fri, 3 Apr 2026 03:57:48 -0400 Subject: [PATCH 10/10] black --- src/gfn/gflownet/trajectory_balance.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/gfn/gflownet/trajectory_balance.py b/src/gfn/gflownet/trajectory_balance.py index ee3e2e4c..2313f5f3 100644 --- a/src/gfn/gflownet/trajectory_balance.py +++ b/src/gfn/gflownet/trajectory_balance.py @@ -242,7 +242,10 @@ def get_scores( Shape ``(N,)`` per-trajectory scores. """ return self._compute_rtb_scores( - env, trajectories, log_rewards=log_rewards, recalculate_all_logprobs=recalculate_all_logprobs + env, + trajectories, + log_rewards=log_rewards, + recalculate_all_logprobs=recalculate_all_logprobs, ) def _compute_rtb_scores(