Simplify state conversions (excludes Crystal environments)

README.md

@@ -1,39 +1,11 @@
-# GFlowNet
+# Private sister repository of gflownet
 
-This repository implements GFlowNets, generative flow networks for probabilistic modelling, on PyTorch. A design guideline behind this implementation is the separation of the logic of the GFlowNet agent and the environments on which the agent can be trained on. In other words, this implementation should allow its extension with new environments without major or any changes to to the agent. Another design guideline is flexibility and modularity. The configuration is handled via the use of [Hydra](https://hydra.cc/docs/intro/).
+This repository (`gflownet-dev`) is private. It is meant to be used to develop research ideas and projects before making them public in the original [alexhernandezgarcia/gflownet](https://github.com/alexhernandezgarcia/gflownet) repository (`gflownet`).
 
-## Installation
+As of October 2023, it is uncertain whether we will stick to this plan in the long term, but the idea is the following:
 
-### pip
+- Develop ideas and projects in `gflownet-dev`.
+- Upon publication or whenever the authors feel comfortable, transfer the relevant code to `gflownet`. 
+- Relevant code improvements and development that does not compromise research projects should be transferred to `gflownet` as early as possible.
 
-```bash
-python -m pip install --upgrade https://github.com/alexhernandezgarcia/gflownet/archive/main.zip
-```
-
-## How to train a GFlowNet model
-
-To train a GFlowNet model with the default configuration, simply run
-
-```bash
-python main.py user.logdir.root=<path/to/log/files/>
-```
-
-Alternatively, you can create a user configuration file in `config/user/<username>.yaml` specifying a `logdir.root` and run
-
-```bash
-python main.py user=<username>
-```
-
-Using Hydra, you can easily specify any variable of the configuration in the command line. For example, to train GFlowNet with the trajectory balance loss, on the continuous torus (`ctorus`) environment and the corresponding proxy:
-
-```bash
-python main.py gflownet=trajectorybalance env=ctorus proxy=torus
-```
-
-The above command will overwrite the `env` and `proxy` default configuration with the configuration files in `config/env/ctorus.yaml` and `config/proxy/torus.yaml` respectively.
-
-Hydra configuration is hierarchical. For instance, a handy variable to change while debugging our code is to avoid logging to wandb. You can do this by setting `logger.do.online=False`.
-
-## Logging to wandb
-
-The repository supports logging of train and evaluation metrics to [wandb.ai](https://wandb.ai), but it is disabled by default. In order to enable it, set the configuration variable `logger.do.online` to `True`.
+This involves extra complexity, so we will re-evaluate or refine this plan after a test period.

gflownet/envs/alaninedipeptide.py

@@ -1,5 +1,5 @@
 from copy import deepcopy
-from typing import List, Tuple
+from typing import List, Tuple, Union
 
 import numpy as np
 import numpy.typing as npt
@@ -40,25 +40,34 @@ def sync_conformer_with_state(self, state: List = None):
             self.conformer.set_torsion_angle(ta, state[idx])
         return self.conformer
 
-    def statetorch2proxy(self, states: TensorType["batch", "state_dim"]) -> npt.NDArray:
+    # TODO: are the conversions to oracle relevant?
+    def states2proxy(
+        self, states: Union[List[List], TensorType["batch", "state_dim"]]
+    ) -> npt.NDArray:
         """
-        Prepares a batch of states in torch "GFlowNet format" for the oracle.
-        """
-        device = states.device
-        if device == torch.device("cpu"):
-            np_states = states.numpy()
-        else:
-            np_states = states.cpu().numpy()
-        return np_states[:, :-1]
-
-    def statebatch2proxy(self, states: List[List]) -> npt.NDArray:
-        """
-        Prepares a batch of states in "GFlowNet format" for the proxy: a tensor where
-        each state is a row of length n_dim with an angle in radians. The n_actions
+        Prepares a batch of states in "environment format" for the proxy: each state is
+        a vector of length n_dim where each value is an angle in radians. The n_actions
         item is removed.
+
+        Important: this method returns a numpy array, unlike in most other
+        environments.
+
+        Args
+        ----
+        states : list or tensor
+            A batch of states in environment format, either as a list of states or as a
+            single tensor.
+
+        Returns
+        -------
+        A numpy array containing all the states in the batch.
         """
-        return np.array(states)[:, :-1]
+        if torch.is_tensor(states[0]):
+            return states.cpu().numpy()[:, :-1]
+        else:
+            return np.array(states)[:, :-1]
 
+    # TODO: need to keep?
     def statetorch2oracle(
         self, states: TensorType["batch", "state_dim"]
     ) -> List[Tuple[npt.NDArray, npt.NDArray]]:
@@ -73,6 +82,7 @@ def statetorch2oracle(
         result = self.statebatch2oracle(np_states)
         return result
 
+    # TODO: need to keep?
     def statebatch2oracle(
         self, states: List[List]
     ) -> List[Tuple[npt.NDArray, npt.NDArray]]: