Merged Box Example Distributions

[josephdviviano]

Sorry for the nasty commit history on this one.

This is an example code that shows how to handle torchgfn in the case that actions can be sampled from two distributions, depending on the state. In this example, when t=0, we sample from one distribution, and when t>0, we sample from a second distribution. Therefore, each batch element of the States tensor must sample from one of two distributions.

DistributionWrapper exposes a .sample() method which sequentially samples from QuarterDisk and QuarterCircleWithExit. These two output spaces are zero-padded so that they are cross-compatible.

BoxPFNeuralNet similarly computes both the S_t=0 parameters, and the S=t>0 parameters, for all batch elements, and then replaces the S=t>0 outputs with the S_t=0 parameters where nessicary.

Note that the number of parameters used for S_t=0 and S=t>0 can differ. We therefore also mask the BoxPFNeuralNet outputs additionally to account for when the smaller number of parameters is utilized.

TODO (In a Follow Up PR):

• Simplify the code. I don't think we need the BoxPFEstimator and BoxPFNeuralNet to be two distinct classes. It would be easier if all the logical existed in the BoxPFEstimator. But we can debate this.
• Ensure the Tabular case works as intended.
• Put the tests in a proper place. But where?

Extras

• Upgraded pyright to 1.1.314.
• Alphabetical sorting and comments in pyproject.toml.
• Identities torch.pi /2 and 2 / torch.pi expressed as globals in box_utils.py.
• Bugfixes in original tests.
• TODOs throughout.

[saleml]

Thanks for the great PR ! I'll merge now, and move the TODOs to a new issue.