Stack(Grid->ContinuousCube) for sampling N continuous D-dimensional vectors

[pieris98]

Hey again @alexhernandezgarcia , @josephdviviano and others,

TL;DR: Can I use Stack(Grid,ContinuousCube) as an environment to sample N*3D-points ($N \times ℝ^3$ vectors)?

Long Version
I've been fiddling with the code to try and build a 3D PointCloud GFlowNet environment as I briefly described in #342 (and #330).

I tried creating a new environment class PointCloudEnv which extends ContinuousCube, but I realized that is probably incorrect because I basically N such ContinuousCubes for $N$ jointly distributed 3D-points (not independent).

The goal is to eventually train a GFN to model the distribution of ModelNet10, but I'd first make it work with a Uniform proxy for testing.

I've also bumped into some errors trying to implement get_logprobs and get_action_space in the process. I tried making a corresponding test_pointcloud.py file to test basic behavior but currently it fails. Trying to run main.py env=pointcloud gives this error:

File "gflownet/gflownet/utils/batch.py", line 209, in set_env
    "state": self.env.state,
AttributeError: 'list' object has no attribute 'state'`

I then stumbled upon Stack used by crystal, which can stack sub-environments together.
I was wondering if it's possible to achieve the desired behavior (N_points * ContinuousCubes with n_dim=3 i.e. 3D) by using a Stack(Grid,ContinuousCube) environment.

1. Is that a plausible way to proceed, or
2. Should I try to implement a new environment from scratch without inheriting from CotninuousCube and instead copying relevant snippets as required?

Thanks all again for all your help in my attempt and excuse any mistake I may have in understanding the concepts!

[pieris98]

P.S. The end-goal would be to train a GFlowNet that does the same job as PointFlow does with Normalizing Flows.
Any pointers as to what may be important or limiting in the case of GFNs for this usecase (dimensionality, convergence, defining a correct policy e.g. MultiVariateGaussians) would be greatly appreciated!