Add Append MDP environment (Set Addition) and Perfect Binary Tree environment (bijection traj -> term states)

[alexandrelarouche]

Hi,

I have implemented two simple gym environments which I find useful when experimenting with GFNs.

First one is simply an append MDP where there is a set of items, and any item can be added to the set at any time. Exit is allowed at any time except at the initial state. An instance may look like the following MDP (without the terminal state to lighten the figure).

The second environment is a Perfect Binary Tree MDP, where there is a bijection between trajectories and terminating states. Action 0 goes left, action 1 goes right. One could easily extend this environment to non-binary trees. We can discuss if it is necessary or interesting for torchgfn to have a non-binary extension environments. An instance of Perfect Binary Tree may look like:

[saleml]

Thank you for the PR @alexandrelarouche
Could you please pre-commit run --all before you commit, then re-push once it succeeds?
For pyright, if you can't solve the issues, you can ignore the lines with a specific comment (that we already used throughout the codevase)

[alexandrelarouche]

Hey @saleml,

My bad, I did run the pre-commit yesterday, but apparently I was negligent in noticing the pyright checks failed. This does open up an avenue of discussion though - it seems strange to require declaring a States type hint in a DiscreteEnv instead of a DiscreteStates. Is that something we could aim to fix or an inherent limitation of pyright's understanding of inheritance?

Thanks for being so pro-active on the repo 🙂 .

[saleml]

Is that something we could aim to fix or an inherent limitation of pyright's understanding of inheritance?
This is definitely something that can and should be fixed.

The first fix should be, indeed, to change all mentions of States to DiscreteStates in the DiscreteEnv class.

There is an open issue (#235) for this type of changes, and we're hoping a pedantic power-user of the library would some day appear and solve it. Perhaps LLMs can do that soon. Otherwise, when/if time permits, I'll do it...

Regarding the two environments you added, are there some published benchmarks, or tests for them?
If you look at the test_scripts.py file, we give some examples of minimal performances to reach for the implemented environments. This should future-proof the environments. I strongly recommend a similar test for each of your environments (of the type: if you run TB with BS 128 using this type of preprocessor, you would get this performance).

Do you think this is something feasible with the current code?

[alexandrelarouche]

Regarding the two environments you added, are there some published benchmarks, or tests for them?
If you look at the test_scripts.py file, we give some examples of minimal performances to reach for the implemented environments. This should future-proof the environments. I strongly recommend a similar test for each of your environments (of the type: if you run TB with BS 128 using this type of preprocessor, you would get this performance).
Do you think this is something feasible with the current code?

As of now, I do not know of specific benchmarks for these specific environments. I use these environment mostly on paper to study GFN behaviors, as they can easily map to the applicative settings I am interested in. I am certain they map to some other applicative scenarios that were studied before, but they offer an abstraction layer above them. As such, I can hardly provide a satisfying test that would show my implementation is correct.

However, as I use them in my own research, I can provide the tests once the papers are out. In the meantime, I understand that this may not be satisfactory for a PR, and we can leave this open if you wish.