Refactoring ideas regarding Sampler

[hyeok9855]

There are two key components, GFlowNet (in src/gfn/gflownet) and Sampler (in src/gfn/samplers.py).

I think Sampler is a redundant class, and we could get rid of it by incorporating the sampling logic inside of the GFlowNet class.

Specifically, this line defines a sampler object only to sample one batch of a trajectory, which can be replaced by a single method in GFlowNet class.

[josephdviviano]

Yes, @saleml and I have had this conversation a few times before. I don't see a major downside of the Sampler becoming a few methods in the base GFlowNet class without much issue.

The only downside I can see is this: if a user needs to subclass the Sampler to implement a new method, it would still be possible with the logic integrated into the GFlowNet, but it would require a little more mental gymnastics. If they were simply subclassing the Sampler, adding some new features, and then passing it to their new GFlowNet the changes would be easier to read / understand.

I'm very open to integrating it.

[josephdviviano]

@hyeok9855 for this one - if you want to propose a PR that handles this merge I will be willing to review.

However I am wary of unnecessarily merging everything into a few giant classes. This is exactly the sort of thing that makes the library harder for other people to extend for their own purposes. What we want is a level of modularity so that it's easy for people to subclass and extend only the required components for a new gflownet method to work.

In this case, having our sampler be distinct from the gflownet is useless because there's only one kind of sampler. But I could imagine in the future having multiple kinds of samplers that are compatible with all of the gflownet modules. Does this make sense to you?

[hyeok9855]

having multiple kinds of samplers that are compatible with all of the gflownet modules. Does this make sense to you?

Yes, that's a very good point. One possible example would be the local search (LS), where we may sample [batch_size * #LS_iterations] trajectories at an epoch.
I think it would be much cleaner to modularize the sampling-with-ls part by, say, LocalSearchSampler rather than do the iterations of LS somewhere else.

[hyeok9855]

Maybe I could implement the LocalSearchSampler, as my first meaningful contribution.

[josephdviviano]

I love this idea!

[hyeok9855]

See #208 for LocalSearchSampler

[hyeok9855]

One thing to discuss regarding sample_trajectories

Currently, we can sample trajectories in two different ways:

1. GFlowNet.sample_trajectories, which defines Sampler inside of the method (both in PFBasedGFlowNet and FMGFlowNet)
2. Sampler.sample_trajectories outside of the GFlowNet obj. (example)

I think this is redundant, so here are options to streamline:

1. Pass Sampler class as an argument of GFlowNet.
2. Remove GFlowNet.sample_trajectories and force to use Sampler.sample_trajectories.

I prefer option 2 because the LocalSearchSampler requires additional arguments (check this).

Please share your ideas!