Add optional custom log_rewards to all GFlowNet losses

src/gfn/gflownet/base.py

@@ -422,6 +422,8 @@ def get_scores(
         trajectories: Trajectories,
         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.
 
@@ -433,6 +435,11 @@ def get_scores(
             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).
 
         Returns:
             A tensor of shape (batch_size,) containing the scores for each trajectory.
@@ -448,11 +455,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)

src/gfn/gflownet/detailed_balance.py

@@ -210,6 +210,8 @@ def get_scores(
         env: Env,
         transitions: Transitions,
         recalculate_all_logprobs: bool = True,
+        *,
+        log_rewards: torch.Tensor | None = None,
     ) -> torch.Tensor:
         r"""Calculates the scores for a batch of transitions.
 
@@ -220,11 +222,31 @@ def get_scores(
             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).
+                **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")
@@ -261,11 +283,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]
@@ -334,6 +358,8 @@ def loss(
         transitions: Transitions,
         recalculate_all_logprobs: bool = True,
         reduction: str = "mean",
+        *,
+        log_rewards: torch.Tensor | None = None,
     ) -> torch.Tensor:
         """Computes the detailed balance loss.
 
@@ -346,6 +372,8 @@ def loss(
             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
@@ -356,6 +384,7 @@ def loss(
         scores = self.get_scores(
             env,
             transitions,
+            log_rewards=log_rewards,
             recalculate_all_logprobs=recalculate_all_logprobs,
         )
         scores = self.loss_fn(scores)

src/gfn/gflownet/flow_matching.py

@@ -281,6 +281,8 @@ def loss(
         states_container: StatesContainer[DiscreteStates],
         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.
 
@@ -295,6 +297,11 @@ def loss(
                 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).
 
         Returns:
             The computed flow matching loss as a tensor. The shape depends on the
@@ -319,10 +326,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

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)
@@ -335,13 +343,20 @@ def get_scores(
         trajectories: Trajectories,
         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).
 
         Returns:
             A tuple (scores, flattening_masks):
@@ -389,6 +404,7 @@ def get_scores(
                 is_terminal_mask,
                 sink_states_mask,
                 i,
+                log_rewards=log_rewards,
             )
 
             flattening_mask = trajectories.terminating_idx.lt(
@@ -561,6 +577,8 @@ def loss(
         trajectories: Trajectories,
         recalculate_all_logprobs: bool = True,
         reduction: str = "mean",
+        *,
+        log_rewards: torch.Tensor | None = None,
     ) -> torch.Tensor:
         """Computes the sub-trajectory balance loss.
 
@@ -572,6 +590,15 @@ def loss(
                 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
+                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).
 
         Returns:
             The computed sub-trajectory balance loss as a tensor. The shape depends on
@@ -598,20 +625,26 @@ 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
             return loss_reduce(self.loss_fn(tb_scores), reduction)
 
         # 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)