Add Graphs as States

[alip67]

Description:
Unlike the current States object that necessitates appending dummy states to batch trajectories of varying lengths, our approach aims to support Trajectories through a nested Batch object representation. The Data class in Torch Geometric represents the graph structure, while the Batch class, which encapsulates batching of Data objects and their efficient indexing, represents the GraphStates object.

The current implementation of Trajectory supports the indexing dimensions: (Num time steps, Num trajectories, State Size). By using a nested Batch of Batch object to represent state Trajectories, the indexing would inherently take the form (Num trajectories, Num timesteps, State size). This approach requires implementing logic within _getitem_() and _setitem_() to internally.

To Do:
Compatibility check with Trajectories, Transition class

[saleml]

Thank you @younik and @alip67 for this important PR. Is there a script we can play with to see the training of the environment you created?

[younik]

Thank you @younik and @alip67 for this important PR. Is there a script we can play with to see the training of the environment you created?

There are still some issues to fix that prevent it from running properly.
I am working to fix them and will post a sample code for using it

[younik]

It looks like GitHub is down, I will rerun CI tomorrow, but they are green locally

[josephdviviano]

First pass at a review. Thanks so much for your amazing work. I am going to mess around with the test cases and environment next to get a better understanding of how things function together. In the meantime, I have some questions.

[josephdviviano]

I think a comment would be worth adding here.

[josephdviviano]

Perhaps we could have a special NoneTensorDict GraphState which acts like None but passes the relevant checks?

[josephdviviano]

This also seems like a much harder constraint.

[younik]

Why? The expected_output_dim is 1 in this class, so it seems the same to me (actually softer as I don't check the dtype).

[josephdviviano]

I disagree that we should fix it here (i.e., in the Estimator), or if we do so, we should implement masking properly in the States class before the release of V2 because this is a clear design pattern violation.

There is no problem adding multiple edges to-from the same node in principle (in say, a multi-attribute graph) but that does make things way more complex, and I think we can safely avoid that complexity here, as single edges to-from the same nodes will cover a lot of AI for Science applications in the near term.

[saleml]

The code runs, thanks for the last change.

A few suggestions/questions:

• If I understand correctly, the goal of the script is to sample ring graphs. I expect that during training, the proportion of sample graphs that are rings gets higher with time. Is it possible to validate that during the training loop? For example, every 10-50-100-whatever iterations, you generate N graphs, and see which ones are rings.
• Could you add short descriptions of the functions you are dealing with? For example, the state_evaluator function seems to be central in the code. It is important to know what it does.
• Could you merge master into this? Note that this might lead to pyright issues. You can solve most, but some, such as gflownet.loss(), you can ignore, using the same comment used to ignore pyright issues from other scripts