BIG PR: Simplification and homegenisation of environments, addition of tests, other clean up changes

[alexhernandezgarcia]

• python -m pytest tests/gflownet/envs/
• black .
• isort .

It's quite a massive PR and it's hard to mention all changes, but here is a non-exhaustive list of things:

• I have simplified the init method of the environments, such that everything that is common to all environments is now in the init of the parent class (base.py) and super() is called at the end of the init.
• Similar thing for the arguments: the arguments of an environment should only be those that are specific to the environment. Everything else goes in the kwargs.
• And methods that were common or very similar for all environments have been unified into the parent class. A good example is reset(). Now it's common for all environments - I think - therefore there is no need to re-implement it.
• All environments now define the self.source state in the init, and this is used by reset() (in the parent class) to initialise the first self.state.
• The end-of-sequence (self.eos) in the environments is now an action with the format of other other actions and not an integer or whatever. This is more consistent in my opinion.
• There is a new method in the parent env class, actions2indices() that is used by the flowmatch loss.
• All proxies can implement a setup() method that receives the environment as argument and sets whatever variables are necessary, like n_dim. setup() is called in the init of the parent env class.
• In general, the proxies have also been cleaned up, and the uniform proxy is now working, which is useful for debugging.
• The buffer class is now in a separate script in utils/buffer.py
• In the grid and the (discrete) torus, extended actions are now handled via the variables max_increment and max_dim_per_action, which specify how many positions (e.g. cells in the grid or discrete angles in the torus) can be done in one action and how many dimensions can be updated in one single action.
• I have re-arranged the order of methods in the environment class.
• I have extended and clarified the docstrings.
• I have added more mypy types - now most methods should have it. If missing, we should add them.
• torus_rounds.py was obsolete and has been removed.#
• I have added a few more assert statements for extra security.
• There are many new tests (although there should be more!) and since many are common to all environments, these are grouped in tests/gflownet/env/common.py
• self.get_actions_space() is now self.get_action_space(); get_max_traj_len is now get_max_traj_length; self.d_action_space is now self.action_space_dim.

Some things remain to be done:

• The aptamers (aptamers.py) environment has not been cleaned up and updated because I have the more ambitious plan of creating a Sequence parent class from which the aptamers will be inherited.
• The plane (plane.py) environment has not been cleaned up and updated because it is not working well yet.
• We need more tests, especially for the continuous environments, and for the GFlowNet agent itself.
• Docstrings need to keep improving.

⚠️ I am not requesting careful reviews of the PR since it is huge, but reviews of things here and there and comments are very welcome. After merging this into continuous, my plan is to merge continuous into main (without further review, obviously) and continue improving things from there.

[nikita-0209]

Regarding the aptamers bullet point, I have created a Sequence parent class from which child classes Aptamers and AMP are derived in the cont_mf branch. I will merge this branch into cont-mf and then create a PR with those changes.

[alexhernandezgarcia]

Regarding the aptamers bullet point, I have created a Sequence parent class from which child classes Aptamers and AMP are derived in the cont_mf branch. I will merge this branch into cont-mf and then create a PR with those changes.

Oh I didn't know you had done this already - that's great, thanks!

[michalkoziarski]

Doesn't have to be a part of this PR, but I think it's time to group envs related to crystal generation into a single directory.

[michalkoziarski]

Could use argument description in docstring (in particular for parent_a).

[AlexandraVolokhova]

I'd like to have it as a method, smth like self.get_source(reset=True) to accommodate scenario when we start from random state

[alexhernandezgarcia]

Yes, to be done when needed 👍