Benchmark

[josephdviviano]

• I've read the .github/CONTRIBUTING.md file
• My code follows the typing guidelines
• I've added appropriate tests
• I've run pre-commit hooks locally

Description

Via the benchmark folder, runs an extensive benchmark of our library against gflownet and gfnx.

The benchmark, after 50 iterations of warm up, times 100 training iterations for each library, for each environment, for each batch size.

Due to the trouble scaling gflownet batch sizes, we had to cap our max batch sizes. However, if this is patched in their project, all of these benchmarks can be re-run.

This required me to add new environments to our library to correctly compare with the other libraries in a few cases:

• BoxCartersian - a continuous environment similar to that in the gflownet library, which is a standard diffusion-style environment.
• The old Box environment is now called "BoxPolar".
• BitSequenceNonAutoregressive - a discrete environment similar to that in gfnx, which is somewhat similar to a discrete diffusion environment (instead of a sequence sampling environment).

Misc changes:

• Turned off annoying None linting.
• Minor performance optimizations throughout the library (easy wins) that arose as a result of the benchmark analysis.
• Some changes to tests / functionality was required to support the transition to BoxCartesian.

[Copilot]

Pull request overview

Adds a new benchmark/ suite to compare torchgfn against gflownet and gfnx, while introducing new/renamed environments (BoxPolar/BoxCartesian and a non-autoregressive BitSequence) plus a set of performance-focused internal optimizations.

Changes:

• Introduces an end-to-end benchmarking framework (runners, scenarios, setup scripts, analysis, result collection).
• Adds new environments/utilities: BoxCartesian, NonAutoregressiveBitSequence, and Cartesian Box estimators/distributions; renames legacy Box to BoxPolar with a compatibility alias.
• Implements several hot-path performance optimizations (States slicing/views + caching, sampler preallocation, conditions propagation, TB weighting numeric consistency).

Reviewed changes

Copilot reviewed 44 out of 47 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
tutorials/examples/train_box_legacy.py	New legacy training script targeting BoxPolar (old semantics).
tutorials/examples/train_box.py	Updates Box training to Cartesian estimators, adds temperature annealing, LR milestone logic, and trajectory plotting.
tutorials/examples/train_bitsequence_non_autoregressive.py	New example script for the non-autoregressive BitSequence environment.
tutorials/examples/test_scripts.py	Updates Box script args/tests to include temperature_start and new JSD targets.
tutorials/examples/output/.gitignore	Ignores generated tutorial outputs (plots/checkpoints).
testing/test_samplers_and_containers.py	Switches sampler/container tests from Box to BoxPolar + imports moved utils.
testing/test_probability_calculations.py	Guards PB transition log-prob computation when no valid next states exist.
testing/test_parametrizations_and_losses.py	Updates Box parametrization tests to Cartesian estimators/modules + seeding.
testing/test_gflownet.py	Updates GFlowNet generic tests to Cartesian Box estimators/modules.
testing/test_environments.py	Reworks Box forward-step tests to validate Cartesian semantics.
src/gfn/utils/modules.py	Fixes/clarifies MLP hidden-layer semantics and adjusts noisy-layer logic + assertions.
src/gfn/states.py	Adds fast _make_view constructor, caches is_initial_state/is_sink_state, and optimizes slicing.
src/gfn/samplers.py	Optimizes trajectory sampling loop (prealloc actions, bypass conditions setter, debug-only storage checks).
src/gfn/gym/tests/test_box_utils.py	Removes old Box util tests (polar-focused).
src/gfn/gym/tests/test_box_polar_utils.py	Adds polar utils tests and additional Cartesian coverage.
src/gfn/gym/tests/test_box_cartesian_utils.py	Adds dedicated test suite for Cartesian Box distribution/estimators.
src/gfn/gym/helpers/box_cartesian_utils.py	New Cartesian Box estimators, MLP heads, and forward/backward distributions.
src/gfn/gym/helpers/bayesian_structure/jsd.py	Replaces deprecated np.float_ with explicit np.float64.
src/gfn/gym/box_cartesian.py	New BoxCartesian env with per-dimension action validation.
src/gfn/gym/box.py	Renames legacy Box to BoxPolar and documents polar validation semantics.
src/gfn/gym/bitSequenceNonAutoregressive.py	New NonAutoregressiveBitSequence discrete environment implementation.
src/gfn/gym/init.py	Exports new envs and sets Box = BoxCartesian for backward compatibility.
src/gfn/gflownet/sub_trajectory_balance.py	Adds TB-equivalent fast path to match TBGFlowNet numeric behavior and avoid cumsum drift.
src/gfn/env.py	Skips conditions setter in internal step/backward_step to reduce overhead.
pyproject.toml	Disables optional (None) handling reports in pyright.
benchmark/setup_benchmark.sh	New setup helper to install benchmark dependencies across platforms.
benchmark/profile_hypergrid.py	New profiling script for HyperGrid iteration breakdown.
benchmark/profile_box.py	New profiling script for Box iteration breakdown.
benchmark/outputs/.gitkeep	Ensures benchmark outputs directory exists in git.
benchmark/outputs/.gitignore	Ignores benchmark result artifacts (CSV/JSON).
benchmark/lib_runners/torchgfn_runner.py	Runner for benchmarking torchgfn across supported envs.
benchmark/lib_runners/gflownet_runner.py	Runner for benchmarking upstream gflownet (Hydra compose + monkeypatches).
benchmark/lib_runners/base.py	Base benchmark dataclasses, runner interface, and result utilities.
benchmark/lib_runners/init.py	Exposes benchmark runner types.
benchmark/gfnx	Adds gfnx as a git submodule pin.
benchmark/gflownet	Adds gflownet as a git submodule pin.
benchmark/collect_results.py	Collects benchmark JSON outputs into a combined CSV.
benchmark/benchmark_libraries.py	Main benchmark entrypoint: scenarios, timing, aggregation, outputs.
benchmark/README.md	Benchmark documentation (setup, scenarios, usage, output format).
benchmark/ANALYSIS.md	Detailed performance analysis and reproducibility notes.
.gitmodules	Declares benchmark submodules.
.gitignore	Ignores .claude/ and fixes .lprof formatting.

Comments suppressed due to low confidence (3)

src/gfn/utils/modules.py:1

• The noisy-layer allocation math can create an extra hidden-to-hidden layer when n_hidden_layers=1 and n_noisy_layers>=2 (because n_extra_hidden=0 but n_noisy_hidden_layers=1). This changes the network depth and parameter count unexpectedly. Fix by ensuring n_noisy_hidden_layers is capped by n_extra_hidden and by accounting separately for whether the input projection and output layer are noisy (e.g., reserve noise for the output first, then allocate any remaining noisy layers to the extra hidden-to-hidden layers only).
  tutorials/examples/train_bitsequence_non_autoregressive.py:1
• This counts terminal-state frequencies with an O(n_states * n_samples) loop and full-tensor comparisons per state. For the documented 'feasible' case (e.g. n_states=4096, n_samples≈409600), this becomes prohibitively large. A faster approach is to encode each terminal state into a unique integer key (e.g., base-n_words positional encoding) and then use torch.bincount (or unique(return_counts=True)) to compute counts in ~O(n_samples).
  tutorials/examples/train_box.py:1
• The comment says the default fires at 33%, 67%, and 100%, but the milestone list uses int(n_iterations / milestone), which typically excludes the final step (e.g., with milestone = n_iterations//3, this yields milestones near 33%, 67%, and ~99% depending on divisibility). Either adjust the milestone construction to match the documented behavior, or update the comment to reflect the actual milestone positions.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

is_action_valid() returns True when all actions are exits, without checking whether exit is valid from the corresponding states (notably: exiting from s0 should be invalid per the Cartesian distribution logic). Consider explicitly validating exit actions before early-returning (e.g., reject actions.is_exit & states.is_initial_state for forward steps; and apply any other exit-availability rules the env intends to enforce).

Suggested change

	if len(non_exit_actions) == 0:
	if len(non_exit_actions) == 0:
	# All actions are exits. For forward steps, exiting from s0 is invalid
	# under the Cartesian semantics, so reject those cases explicitly.
	if not backward:
	if torch.any(states.is_initial_state & actions.is_exit):
	return False

[Copilot]

BoxCartesianDistribution.log_prob() clamps r into (eps, 1-eps) without enforcing the action support, so out-of-range actions (e.g., < delta for non-s0, or > 1 - state) can receive finite log-prob instead of -inf. This breaks the expected 'probability 0 outside support' invariant and can hide invalid-action issues. Suggest adding an explicit validity mask based on state (s0 vs non-s0) and setting log_probs to -inf for non-exit actions that violate per-dimension bounds.

[Copilot]

This PR introduces a new public environment (NonAutoregressiveBitSequence) but doesn't add a corresponding test module validating core invariants (forward/backward mask correctness, step/backward_step roundtrip, exit-only-at-terminal behavior, and reward monotonicity vs Hamming distance). Adding targeted pytest coverage would help prevent regressions in action encoding/masking and terminal/exit semantics.