add support for including Graphs as States of GFlowNet

[ashdtu]

Describe your changes

Type: Feature
Add support for including Graphs as States of GFlowNet
The Graph structure is represented via the Data class of Torch Geometric. The GraphStates object is represented via the Batch class which encapsulates batching of Data objects and their efficient indexing. As opposed to the existing States object which requires appending dummy states for batching different length Trajectories together, we seek to support Trajectories by representing it as nested Batch object.

The existing implementation of Trajectory supports the following indexing dimensions: (Num time steps, Num trajectories, State Size). The nested Batch of Batch object representing state Trajectories would naturally indexing of the form (Num trajectories, Num timesteps, State size), and would need to implement logic for flipping the indexing dimensions internally in _getitem_() and _setitem_().

Issue Number

#153

Testing

• Unit tests for helper functions of GraphState class: initialization, appending, batching etc.
• Compatibility check with Trajectories, Transition class
• Unit testing with FlowMatching : accessing GraphStates directly in loss function
• Unit testing with TrajectoryBalance: accessing GraphStates via trajectory class

[LicoriceLin]

Hi, is there any further plan on implementing this feature? What's the bugs so far? Is there anything I could help out?

And if I just want to hack it, can I simply copy the GraphStates class as a temporary solution?

[josephdviviano]

Hi @LicoriceLin yes there is a plan and it's now a top priority. A few of us will be looking at adding this feature starting next week and we would love your help and or feedback on the implementation to ensure we're able to support your needs. CC @younik