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 1/2] 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 2/2] 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.