Detachment of env and proxy; avoiding copies of the environments

[alexhernandezgarcia]

This PR refactors several important aspects of the code base:

• Avoid unnecessary copies of environments to sample trajectories.
• Detachment of environment and proxy
• Dropping the convention of range of values in the proxies

Copies of environments

Each trajectory in a GFlowNet batch is constructed by evolving one environment, either forwards or backwards. So far, to construct each trajectory, a reference environment was copied and reset. This PR replaces this copying operation by storing an "environment maker" as an attribute of the agent (as though it was just the environment class) and instead creating a new environment instance for each trajectory.

Detachment of environment and proxy

This addresses Issue #288.

So far, the proxy has been part of GFlowNet environments as an attribute. The rationale was to be able to compute the reward of a trajectory directly from the environment, for example with env.reward(). While this is a nice feature, it comes at a cost and it has prevented extending the flexibility to use different proxy-to-reward functions, a better handling of the computation of log-rewards for better numerical stability, as well the simple use of alternative baselines to GFlowNets.

In this PR, the proxy and the environments have been detached and become (almost fully) independent of each other. Together with this change, I have also tried to improve the implementation for handling the conversion of proxy values to rewards or log rewards, to make it for efficient, numerically more stable and more flexible to extend.

These are some of the new features:

• Several proxy-to-reward functions are transparently supported: identity, absolute value, power (to β), exponential (Boltzmann), shift (sum of β) and product (with β).
• All this is now implemented in the base Proxy class.
• One can easily create a new custom proxy-to-reward function in their custom proxy class and use it by simply setting it as a callable to in self.reward_function.
• Since all losses are computed in log domain, now the log-rewards can directly be computed from the proxy values without first computing rewards and then taking the log. This may be particularly important for numerical stability when using exponential proxy-to-reward functions.

Dropping convention of range of values in the proxies

So far, we (sort of) followed a convention by which lower proxy values (and negative) were better. This has created a lot of confusion because not all proxies make much sense in this convention.

With this PR, this convention is dropped altogether and the guiding principle now is that the user is free to create any proxy with any range of values and any direction (maximisation or minimisation), but they are responsible to use an appropriate proxy-to-reward function. This is now much easier given the changes described in the previous section.

To-do list

Some of these things may be left for a separate PR.

• Sanity check experiments
• IMPORTANT: fix FL loss with exponential proxy2reward (see Tetris) - currently optimises opposite direction
• Fix sanity check experiments with tetris plotting
• Adapt methods in environments (e.g. sample_from_reward(), compute_train_energy_proxy_and_rewards()...) that still use the old methods in the base env.
• Consider getting rid of self.reward_norm and co. in the base env.
• Fix broken common tests with minimal runs.
• Update all configs...
• Avoid re-computing proxy values for Buffer This is left for a future PR.
• Update outputs of all proxies: the output does not need to be negative anymore, then we can make identity the default.
• Write doc in Proxy base about reward_func and rest of attributes.

[alexhernandezgarcia]

Ok, to me this is now ready to merge. @carriepl @AlexandraVolokhova I'm tagging you as reviewers in case you want to take a quick look and to let you know. Feel free to go ahead and approve and merge unless you have any concerns. Sanity checks look perfect: https://wandb.ai/alexhg/gfn_sanity_checks (the issue with FL was that the sanity checks config files was incorrect for the new proxy - fixed in 361cdc0)

[alexhernandezgarcia]

@ginihumer I am tagging you here because this will be merged soon into the main branch and then into the activelearning branch, which will force making changes (for the better) on the active learning repo. No need to check the code details, but it's worth reading the description at least to understand what will change. In a nutshell, that proxies don't need to be negative and the lower the better, but can be anything and can be controlled by reward_function. Also, that the proxy (e.g. acquisition function) won't be part of the environment any more.

[carriepl]

There's a huge amount of work in this PR! Overall, it looks pretty good to me.

[carriepl]

This line is mysterious to me. Not sure if I'll understand it later when I read the remaining changes so... just leaving a comment here as a reminder.

[alexhernandezgarcia]

Oh dear, good catch! This is pretty wrong and it's gonna need a bunch of changes that will delay things a bit...

[alexhernandezgarcia]

How about this? a1462f1

[alexhernandezgarcia]

@carriepl Thank you so much for the review! Would you mind taking a quick look at my comments and resolving yours if the comment is addressed. There should be one remaining issue only, if I am not mistaken.

[carriepl]

@carriepl Thank you so much for the review! Would you mind taking a quick look at my comments and resolving yours if the comment is addressed. There should be one remaining issue only, if I am not mistaken.

I just looked. It all looks great. Only that one issue remaining.

[alexhernandezgarcia]

Yes this is nice! I use it more and more, for example in the tests of the new Crystal subenvs.

On 21 May 2024 12:59:01 GMT-04:00, carriepl ***@***.***> wrote: @carriepl commented on this pull request. > + 1.0, + 8.1031e03, + ], + [-11, -2, -1.5, -1.1, -1, -0.9, -0.5, 0, 9], + ), + ], +) +def test_reward_function_callable__behaves_as_expected( + proxy_callable, + reward_function, + logreward_function, + proxy_values, + rewards_exp, + logrewards_exp, +): + proxy = proxy_callable Ok, I wasn't aware that Pytest fixtures could implicitely use the test parameters. Good to know! -- Reply to this email directly or view it on GitHub: #299 (comment) You are receiving this because you authored the thread. Message ID: ***@***.***>

-- Sent from /e/ Mail.

[carriepl]

Looks good to me. Many thanks!