Repair states2proxy in Composition

alexhernandezgarcia · Dec 9, 2023 · c5763ea · c5763ea
1 parent b07e258
commit c5763ea
Showing 1 changed file with 12 additions and 3 deletions.
diff --git a/gflownet/envs/crystals/composition.py b/gflownet/envs/crystals/composition.py
@@ -406,8 +406,10 @@ def states2proxy(
         self, states: Union[List[List], TensorType["batch", "state_dim"]]
     ) -> TensorType["batch", "state_proxy_dim"]:
         """
-        Prepares a batch of states in "environment format" for the proxy: simply
-        returns the states as are with dtype long.
+        Prepares a batch of states in "environment format" for the proxy: The output is
+        a tensor of dtype long with N_ELEMENTS_ORACLE + 1 columns, where the positions
+        of self.elements are filled with the number of atoms of each element in the
+        state.
 
         Args
         ----
@@ -419,7 +421,14 @@ def states2proxy(
         -------
         A tensor containing all the states in the batch.
         """
-        return tlong(states, device=self.device)
+        states = tlong(states, device=self.device)
+        states_proxy = torch.zeros(
+            (states.shape[0], N_ELEMENTS_ORACLE + 1),
+            device=self.device,
+            dtype=torch.long,
+        )
+        states_proxy[:, tlong(self.elements, device=self.device)] = states
+        return states_proxy
 
     def state2readable(self, state=None):
         """