trajectories and transitions now initalize log_rewards correctly, so …

…extend always works, which means we no longer need the extra condition (but we are leaving in the sanity check for now).

src/gfn/containers/trajectories.py

@@ -90,7 +90,11 @@ def __init__(
             if when_is_done is not None
             else torch.full(size=(0,), fill_value=-1, dtype=torch.long)
         )
-        self._log_rewards = log_rewards
+        self._log_rewards = (
+            log_rewards
+            if log_rewards is not None
+            else torch.full(size=(0,), fill_value=0, dtype=torch.float)
+        )
         self.log_probs = (
             log_probs
             if log_probs is not None
@@ -246,10 +250,7 @@ def extend(self, other: Trajectories) -> None:
                 (self._log_rewards, other._log_rewards),
                 dim=0,
             )
-        # If the trajectories object does not yet have `log_rewards` assigned but the
-        # external trajectory has log_rewards, simply assign them over.
-        elif self._log_rewards is None and other._log_rewards is not None:
-            self._log_rewards = other._log_rewards
+        # Will not be None if object is initialized as empty.
         else:
             self._log_rewards = None
 

src/gfn/containers/transitions.py

@@ -90,7 +90,7 @@ def __init__(
             len(self.next_states.batch_shape) == 1
             and self.states.batch_shape == self.next_states.batch_shape
         )
-        self._log_rewards = log_rewards
+        self._log_rewards = log_rewards if log_rewards is not None else torch.zeros(0)
         self.log_probs = log_probs if log_probs is not None else torch.zeros(0)
 
     @property
@@ -208,6 +208,8 @@ def extend(self, other: Transitions) -> None:
         self.actions.extend(other.actions)
         self.is_done = torch.cat((self.is_done, other.is_done), dim=0)
         self.next_states.extend(other.next_states)
+
+        # Concatenate log_rewards of the trajectories.
         if self._log_rewards is not None and other._log_rewards is not None:
             self._log_rewards = torch.cat(
                 (self._log_rewards, other._log_rewards), dim=0