Scrabble env

config/env/scrabble.yaml

@@ -0,0 +1,15 @@
+defaults:
+  - base
+
+_target_: gflownet.envs.scrabble.Scrabble
+
+id: scrabble
+# Buffer
+buffer:
+  data_path: null
+  train: null
+  test:
+    type: uniform
+    n: 10
+    output_csv: scrabble_test.csv
+    output_pkl: scrabble_test.pkl

config/proxy/scrabble.yaml

@@ -0,0 +1,3 @@
+_target_: gflownet.proxy.scrabble.ScrabbleScorer
+
+vocabulary_check: False

config/user/alex.yaml

@@ -1,5 +1,4 @@
 logdir:
-  root: /network/scratch/h/hernanga/logs/gflownet
+  root: /home/alex/logs/gflownet
 data:
-  root: /home/mila/h/hernanga/gflownet/data
-  alanine_dipeptide: /home/mila/h/hernanga/gflownet/data/alanine_dipeptide_conformers_1.npy
+  root: /home/alex/datasets

gflownet/envs/base.py

@@ -694,7 +694,16 @@ def state2proxy(self, state: Union[List, TensorType["state_dim"]] = None):
             A state
         """
         state = self._get_state(state)
-        return torch.squeeze(self.states2proxy([state]), dim=0)
+        state_proxy = self.states2proxy([state])
+        if isinstance(state_proxy, list):
+            return state_proxy[0]
+        elif torch.is_tensor(state_proxy):
+            return torch.squeeze(state_proxy, dim=0)
+        else:
+            raise NotImplementedError(
+                "The output of states2proxy must be either a list or a tensor. "
+                f"Got {type(state_proxy)}."
+            )
 
     def states2policy(
         self, states: Union[List, TensorType["batch", "state_dim"]]