estimator_outputs can be passed around

src/gfn/gflownet/trajectory_balance.py

@@ -41,7 +41,12 @@ def __init__(
         self.logZ = nn.Parameter(torch.tensor(init_logZ))
         self.log_reward_clip_min = log_reward_clip_min
 
-    def loss(self, env: Env, trajectories: Trajectories) -> TT[0, float]:
+    def loss(
+            self,
+            env: Env,
+            trajectories: Trajectories,
+            estimator_outputs: torch.Tensor = None,
+    ) -> TT[0, float]:
         """Trajectory balance loss.
 
         The trajectory balance loss is described in 2.3 of
@@ -51,7 +56,7 @@ def loss(self, env: Env, trajectories: Trajectories) -> TT[0, float]:
             ValueError: if the loss is NaN.
         """
         del env  # unused
-        _, _, scores = self.get_trajectories_scores(trajectories)
+        _, _, scores = self.get_trajectories_scores(trajectories, estimator_outputs)
         loss = (scores + self.logZ).pow(2).mean()
         if torch.isnan(loss):
             raise ValueError("loss is nan")
@@ -80,14 +85,19 @@ def __init__(
 
         self.log_reward_clip_min = log_reward_clip_min
 
-    def loss(self, env: Env, trajectories: Trajectories) -> TT[0, float]:
+    def loss(
+            self,
+            env: Env,
+            trajectories: Trajectories,
+            estimator_outputs: torch.Tensor = None,
+        ) -> TT[0, float]:
         """Log Partition Variance loss.
 
         This method is described in section 3.2 of
         [ROBUST SCHEDULING WITH GFLOWNETS](https://arxiv.org/abs/2302.05446))
         """
         del env  # unused
-        _, _, scores = self.get_trajectories_scores(trajectories)
+        _, _, scores = self.get_trajectories_scores(trajectories, estimator_outputs)
         loss = (scores - scores.mean()).pow(2).mean()
         if torch.isnan(loss):
             raise ValueError("loss is NaN.")