Add optional custom log_rewards to all GFlowNet losses

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)

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)

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

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,20 +619,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)