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/main.yaml

@@ -3,6 +3,7 @@ defaults:
   - env: grid
   - gflownet: flowmatch
   - policy: mlp_${gflownet}
+  - state_flow: null
   - proxy: corners
   - logger: wandb
   - user: alex

config/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

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",
@@ -76,9 +77,15 @@ def __init__(
         if optimizer.loss in ["flowmatch", "flowmatching"]:
             self.loss = "flowmatch"
             self.logZ = None
+            self.non_terminal_rewards = False
         elif optimizer.loss in ["trajectorybalance", "tb"]:
             self.loss = "trajectorybalance"
             self.logZ = nn.Parameter(torch.ones(optimizer.z_dim) * 150.0 / 64)
+            self.non_terminal_rewards = False
+        elif optimizer.loss in ["forwardlooking", "fl"]:
+            self.loss = "forwardlooking"
+            self.logZ = None
+            self.non_terminal_rewards = True
         else:
             print("Unkown loss. Using flowmatch as default")
             self.loss = "flowmatch"
@@ -121,7 +128,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 +141,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 +154,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 +195,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 nn 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 in {self.loss} loss.")
+            parameters += list(self.state_flow.model.parameters())
+        return parameters
 
     def sample_actions(
         self,
@@ -405,12 +424,22 @@ def sample_batch(
             "actions_envs": 0.0,
         }
         t0_all = time.time()
-        batch = Batch(env=self.env, device=self.device, float_type=self.float)
+        batch = Batch(
+            env=self.env,
+            device=self.device,
+            float_type=self.float,
+            non_terminal_rewards=self.non_terminal_rewards,
+        )
 
         # ON-POLICY FORWARD trajectories
         t0_forward = time.time()
         envs = [self.env.copy().reset(idx) for idx in range(n_forward)]
-        batch_forward = Batch(env=self.env, device=self.device, float_type=self.float)
+        batch_forward = Batch(
+            env=self.env,
+            device=self.device,
+            float_type=self.float,
+            non_terminal_rewards=self.non_terminal_rewards,
+        )
         while envs:
             # Sample actions
             t0_a_envs = time.time()
@@ -432,7 +461,12 @@ def sample_batch(
         # TRAIN BACKWARD trajectories
         t0_train = time.time()
         envs = [self.env.copy().reset(idx) for idx in range(n_train)]
-        batch_train = Batch(env=self.env, device=self.device, float_type=self.float)
+        batch_train = Batch(
+            env=self.env,
+            device=self.device,
+            float_type=self.float,
+            non_terminal_rewards=self.non_terminal_rewards,
+        )
         if n_train > 0 and self.buffer.train_pkl is not None:
             with open(self.buffer.train_pkl, "rb") as f:
                 dict_tr = pickle.load(f)
@@ -463,7 +497,12 @@ def sample_batch(
 
         # REPLAY BACKWARD trajectories
         t0_replay = time.time()
-        batch_replay = Batch(env=self.env, device=self.device, float_type=self.float)
+        batch_replay = Batch(
+            env=self.env,
+            device=self.device,
+            float_type=self.float,
+            non_terminal_rewards=self.non_terminal_rewards,
+        )
         if n_replay > 0 and self.buffer.replay_pkl is not None:
             with open(self.buffer.replay_pkl, "rb") as f:
                 dict_replay = pickle.load(f)
@@ -663,6 +702,79 @@ 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)
+
+        assert batch.non_terminal_rewards
+        rewards_states = batch.get_rewards()
+        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,
@@ -753,7 +865,12 @@ def estimate_logprobs_data(
         end_batch = min(batch_size, n_states)
         pbar = tqdm(total=n_states)
         while init_batch < n_states:
-            batch = Batch(env=self.env, device=self.device, float_type=self.float)
+            batch = Batch(
+                env=self.env,
+                device=self.device,
+                float_type=self.float,
+                non_terminal_rewards=self.non_terminal_rewards,
+            )
             # Create an environment for each data point and trajectory and set the state
             envs = []
             for state_idx in range(init_batch, end_batch):
@@ -852,7 +969,12 @@ def train(self):
                 self.logger.log_metrics(metrics, use_context=self.use_context, step=it)
                 self.logger.log_summary(summary)
             t0_iter = time.time()
-            batch = Batch(env=self.env, device=self.device, float_type=self.float)
+            batch = Batch(
+                env=self.env,
+                device=self.device,
+                float_type=self.float,
+                non_terminal_rewards=self.non_terminal_rewards,
+            )
             for j in range(self.sttr):
                 sub_batch, times = self.sample_batch(
                     n_forward=self.batch_size.forward,
@@ -869,6 +991,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 +1056,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 +1087,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()
@@ -1136,7 +1264,12 @@ def test_top_k(self, it, progress=False, gfn_states=None, random_states=None):
         print()
         if not gfn_states:
             # sample states from the current gfn
-            batch = Batch(env=self.env, device=self.device, float_type=self.float)
+            batch = Batch(
+                env=self.env,
+                device=self.device,
+                float_type=self.float,
+                non_terminal_rewards=self.non_terminal_rewards,
+            )
             self.random_action_prob = 0
             t = time.time()
             print("Sampling from GFN...", end="\r")
@@ -1159,7 +1292,12 @@ def test_top_k(self, it, progress=False, gfn_states=None, random_states=None):
         if do_random:
             # sample random states from uniform actions
             if not random_states:
-                batch = Batch(env=self.env, device=self.device, float_type=self.float)
+                batch = Batch(
+                    env=self.env,
+                    device=self.device,
+                    float_type=self.float,
+                    non_terminal_rewards=self.non_terminal_rewards,
+                )
                 self.random_action_prob = 1.0
                 print("[test_top_k] Sampling at random...", end="\r")
                 for b in batch_with_rest(