Rethinking sampling

[josephdviviano]

This PR is a hodgepodge of a few tweaks, bugfixes, and investigations related to the sampling logic, including a new simple continuous example.

• estimator_outputs are now re-used when sampling off policy.
  • This is accomplished currently using padding. In a future PR this will be handled using nested tensors. I will likely wait until their API solidifies (it will apparently change in the near future).
• policy_kwargs are passed around properly to do off policy sampling, hopefully in a way that remains generic.
• TODO:s added around based on my observations.
• log_reward_min_clamp now off by default and only defined in the gflownets, NOT the environment.
• Other minor tweaks I made while chasing a ghost relating to magically changing numbers (I now believe this to be due to faulty RAM on my laptop).
• API change: the user must now specify often whether a gflownet or sampler is running off_policy or not. This is for efficiency. If sampling happens off policy, we save estimator outputs (because we assume we will need them later, to evaluate log probabilities of actions under the policy). If it's done on_policy, we calculate the log-rewards during the forward pass.

[saleml]

I'll wait for the other PR to merge so that diffs are accurate here

[josephdviviano]

Sorry in advance for the large PR - feel free to be critical ...

[josephdviviano]

Thanks for the feedback!

I really like your idea of associating each TODO with an issue. That would also make it easier to go fix the thing (you just search for the issue number in the code).

I can do this in a follow up PR!

Sorry, I knew I was being naughty when I submitted this monster PR. It essentially was a grab bag of things I tried, while trying to get the library to play ball for the gflownet workshop, and it would have been really annoying to go split it out into various PRs post-hoc. I figured it wasn't too bad because I was the only one working on it but I agree this is horrible practice and not conducive to collaboration.

I understand your desire for strong typing on the policy_kwargs but I'm worried it will add a lot of developer overhead. We should keep this in mind. In general, I don't want researchers to have to think too hard about software engineering stuff when using this library - we should figure out the minimal set of good engineering practices that are researcher friendly. To be frank, the variance in engineering skills in research is enormous, because the pipeline does not select much for engineering ability, and I think this library will have the greatest impact if we can make it accessible to as many of those people as possible.

And just to clarify the intent of the PR: I was addressing multiple points of feedback:

• It's inefficient to do two forward passes on a neural network when one would suffice.
• The intention to train off or on policy was too often implicit. It wasn't obvious to new users how to sample off policy. So now everything is very explicit (and often required to be passed by the user).
• The copy stuff is arising from discussions with a collaborator at Intel who did some profiling of the library. I think that piece is far from over - but the stuff I did was mostly trying to track down that bug with the slightly changing values.

[josephdviviano]

I'll wait for @saleml (who is defending next week so is likely distracted at the minute) to merge. No rush!

[saleml]

On it !

[josephdviviano]

Hey Salem -- the indexing changes I implemented to fix this

https://github.com/GFNOrg/torchgfn/pull/147#discussion_r1430087521

have broken the tests -- I'm working on that now. I opened a can of worms here! There's likely an elegant solution.

[josephdviviano]

Finally I revered the change for that failing test. I'm not sure there's a better solution that isn't extremely complicated.

I'd like to get this PR merged but we can happily revisit this issue in a future much smaller PR.

[josephdviviano]

@saleml would love you to check this before I merge :)

[saleml]

This is some great work! Thank you Joseph for this very important PR.

I have read your replies to my comments, and seen that the tests pass. I think this can be merged as is.

Small nitpick: For this comment: https://github.com/GFNOrg/torchgfn/pull/147#discussion_r1430099166, do you think we should add the argument in the abstract function?

[josephdviviano]

Small nitpick: For this comment: https://github.com/GFNOrg/torchgfn/pull/147#discussion_r1430099166, do you think we should add the argument in the abstract function?

I added it!