Forward-looking (FL) and Detailed Balance (DB) losses

config/experiments/simple_tetris.yaml

@@ -0,0 +1,47 @@
+# @package _global_
+
+defaults:
+  - override /env: tetris
+  - override /gflownet: trajectorybalance
+  - override /policy: mlp
+  - override /proxy: tetris
+  - override /logger: wandb
+
+env:
+  reward_func: boltzmann
+  reward_beta: 10.0
+  width: 4
+  height: 4
+  pieces: ["I", "O", "J", "L", "T"]
+  rotations: [0, 90, 180, 270]
+  buffer:
+    # replay_capacity: 0
+    test:
+      type: random
+      output_csv: simple_tetris_val.csv
+      output_pkl: simple_tetris_val.pkl
+      n: 100
+
+gflownet:
+  random_action_prob: 0.3
+  optimizer:
+    n_train_steps: 10000
+    lr_z_mult: 100
+    lr: 0.0001
+
+policy:
+  forward:
+    type: mlp
+    n_hid: 128
+    n_layers: 5
+
+  backward:
+    shared_weights: True
+    checkpoint: null
+    reload_ckpt: False
+
+device: cpu
+logger:
+  do:
+    online: True
+  project_name: simple_tetris

config/gflownet/forwardlooking.yaml

@@ -0,0 +1,9 @@
+defaults:
+  - gflownet
+  - state_flow: mlp
+
+optimizer:
+  loss: forwardlooking
+  lr: 0.0001
+  lr_decay_period: 1000000
+  lr_decay_gamma: 0.5

config/gflownet/gflownet.yaml

@@ -34,6 +34,8 @@ optimizer:
   # From original implementation
   bootstrap_tau: 0.0
   clip_grad_norm: 0.0
+# State flow modelling
+state_flow: null
 # If True, compute rewards in batches
 batch_reward: True
 # Force zero probability of sampling invalid actions

config/gflownet/state_flow/mlp.yaml

@@ -0,0 +1,9 @@
+_target_: gflownet.policy.state_flow.StateFlow
+
+config:
+  type: mlp
+  n_hid: 128
+  n_layers: 2
+  checkpoint: null
+  reload_ckpt: False
+  shared_weights: False

config/policy/mlp_forwardlooking.yaml

@@ -0,0 +1,7 @@
+defaults:
+  - mlp
+
+backward:
+  shared_weights: True
+  checkpoint: null
+  reload_ckpt: False

gflownet/envs/base.py

@@ -780,13 +780,13 @@ def traj2readable(self, traj=None):
         """
         return str(traj).replace("(", "[").replace(")", "]").replace(",", "")
 
-    def reward(self, state=None, done=None):
+    def reward(self, state=None, done=None, do_non_terminating=False):
         """
         Computes the reward of a state
         """
         state = self._get_state(state)
         done = self._get_done(done)
-        if done is False:
+        if done is False and do_non_terminating is False:
             return tfloat(0.0, float_type=self.float, device=self.device)
         return self.proxy2reward(self.proxy(self.state2proxy(state))[0])
 

gflownet/gflownet.py

@@ -50,6 +50,7 @@ def __init__(
         logger,
         num_empirical_loss,
         oracle,
+        state_flow=None,
         active_learning=False,
         sample_only=False,
         replay_sampling="permutation",
@@ -79,6 +80,9 @@ def __init__(
         elif optimizer.loss in ["trajectorybalance", "tb"]:
             self.loss = "trajectorybalance"
             self.logZ = nn.Parameter(torch.ones(optimizer.z_dim) * 150.0 / 64)
+        elif optimizer.loss in ["forwardlooking", "fl"]:
+            self.loss = "forwardlooking"
+            self.logZ = None
         else:
             print("Unkown loss. Using flowmatch as default")
             self.loss = "flowmatch"
@@ -121,7 +125,8 @@ def __init__(
             print(f"\tStd score: {self.buffer.test['energies'].std()}")
             print(f"\tMin score: {self.buffer.test['energies'].min()}")
             print(f"\tMax score: {self.buffer.test['energies'].max()}")
-        # Policy models
+
+        # Models
         self.forward_policy = forward_policy
         if self.forward_policy.checkpoint is not None:
             self.logger.set_forward_policy_ckpt_path(self.forward_policy.checkpoint)
@@ -133,6 +138,7 @@ def __init__(
                 print("Reloaded GFN forward policy model Checkpoint")
         else:
             self.logger.set_forward_policy_ckpt_path(None)
+
         self.backward_policy = backward_policy
         self.logger.set_backward_policy_ckpt_path(None)
         if self.backward_policy.checkpoint is not None:
@@ -145,6 +151,14 @@ def __init__(
                 print("Reloaded GFN backward policy model Checkpoint")
         else:
             self.logger.set_backward_policy_ckpt_path(None)
+
+        self.state_flow = state_flow
+        if self.state_flow is not None and self.state_flow.checkpoint is not None:
+            self.logger.set_state_flow_ckpt_path(self.state_flow.checkpoint)
+            # TODO: add the logic and conditions to reload a model
+        else:
+            self.logger.set_state_flow_ckpt_path(None)
+
         # Optimizer
         if self.forward_policy.is_model:
             self.target = copy.deepcopy(self.forward_policy.model)
@@ -178,14 +192,16 @@ def __init__(
         self.nll_tt = 0.0
 
     def parameters(self):
-        if self.backward_policy.is_model is False:
-            return list(self.forward_policy.model.parameters())
-        elif self.loss == "trajectorybalance":
-            return list(self.forward_policy.model.parameters()) + list(
-                self.backward_policy.model.parameters()
-            )
-        else:
-            raise ValueError("Backward Policy cannot be a nn in flowmatch.")
+        parameters = list(self.forward_policy.model.parameters())
+        if self.backward_policy.is_model:
+            if self.loss == "flowmatch":
+                raise ValueError("Backward Policy cannot be a model in flowmatch.")
+            parameters += list(self.backward_policy.model.parameters())
+        if self.state_flow is not None:
+            if self.loss != "forwardlooking":
+                raise ValueError(f"State flow cannot be trained with {self.loss} loss.")
+            parameters += list(self.state_flow.model.parameters())
+        return parameters
 
     def sample_actions(
         self,
@@ -663,6 +679,78 @@ def trajectorybalance_loss(self, it, batch):
         )
         return loss, loss, loss
 
+    def forwardlooking_loss(self, it, batch):
+        """
+        Computes the Forward-Looking GFlowNet loss of a batch
+        Reference : https://arxiv.org/pdf/2302.01687.pdf
+
+        Args
+        ----
+        it : int
+            Iteration
+
+        batch : Batch
+            A batch of data, containing all the states in the trajectories.
+
+
+        Returns
+        -------
+        loss : float
+
+        term_loss : float
+            Loss of the terminal nodes only
+
+        nonterm_loss : float
+            Loss of the intermediate nodes only
+        """
+
+        assert batch.is_valid()
+        # Get necessary tensors from batch
+        states_policy = batch.get_states(policy=True)
+        states = batch.get_states(policy=False)
+        actions = batch.get_actions()
+        parents_policy = batch.get_parents(policy=True)
+        parents = batch.get_parents(policy=False)
+        traj_indices = batch.get_trajectory_indices(consecutive=True)
+        done = batch.get_done()
+
+        masks_b = batch.get_masks_backward()
+        policy_output_b = self.backward_policy(states_policy)
+        logprobs_bkw = self.env.get_logprobs(
+            policy_output_b, actions, masks_b, states, is_backward=True
+        )
+        masks_f = batch.get_masks_forward(of_parents=True)
+        policy_output_f = self.forward_policy(parents_policy)
+        logprobs_fwd = self.env.get_logprobs(
+            policy_output_f, actions, masks_f, parents, is_backward=False
+        )
+
+        states_log_flflow = self.state_flow(states_policy)
+        # forward-looking flow is 1 in the terminal states
+        states_log_flflow[done.eq(1)] = 0.0
+        # Can be optimised by reusing states_log_flflow and batch.get_parent_indices
+        parents_log_flflow = self.state_flow(parents_policy)
+
+        rewards_states = batch.get_rewards(do_non_terminating=True)
+        rewards_parents = batch.get_rewards_parents()
+        energies_states = -torch.log(rewards_states)
+        energies_parents = -torch.log(rewards_parents)
+
+        per_node_loss = (
+            parents_log_flflow
+            - states_log_flflow
+            + logprobs_fwd
+            - logprobs_bkw
+            + energies_states
+            - energies_parents
+        ).pow(2)
+
+        term_loss = per_node_loss[done].mean()
+        nonterm_loss = per_node_loss[~done].mean()
+        loss = per_node_loss.mean()
+
+        return loss, term_loss, nonterm_loss
+
     @torch.no_grad()
     def estimate_logprobs_data(
         self,
@@ -869,6 +957,8 @@ def train(self):
                     losses = self.trajectorybalance_loss(
                         it * self.ttsr + j, batch
                     )  # returns (opt loss, *metrics)
+                elif self.loss == "forwardlooking":
+                    losses = self.forwardlooking_loss(it * self.ttsr + j, batch)
                 else:
                     print("Unknown loss!")
                 # TODO: deal with this in a better way
@@ -932,7 +1022,9 @@ def train(self):
             times.update({"log": t1_log - t0_log})
             # Save intermediate models
             t0_model = time.time()
-            self.logger.save_models(self.forward_policy, self.backward_policy, step=it)
+            self.logger.save_models(
+                self.forward_policy, self.backward_policy, self.state_flow, step=it
+            )
             t1_model = time.time()
             times.update({"save_interim_model": t1_model - t0_model})
 
@@ -961,7 +1053,9 @@ def train(self):
             self.logger.log_time(times, use_context=self.use_context)
 
         # Save final model
-        self.logger.save_models(self.forward_policy, self.backward_policy, final=True)
+        self.logger.save_models(
+            self.forward_policy, self.backward_policy, self.state_flow, final=True
+        )
         # Close logger
         if self.use_context is False:
             self.logger.end()