BrainBlock

BrainBlock is a reproducible reinforcement-learning project for an 8x5 tetromino packing environment. It includes a Gymnasium-compatible environment, PyTorch masked Double DQN training, deterministic held-out evaluation, baseline policies, exact-solver examples, queue-disjointness auditing, plots, and visual rollout tools.

Features

• Gymnasium-compatible BrainBlock environment with a fixed tetromino inventory.
• Full 8 * 8 * 5 = 320 orientation/anchor discrete action space.
• Action-masked Double DQN implemented directly in PyTorch.
• Sparse, dense, feasible-shaping, and no-expert ablation configurations.
• Deterministic train, validation, expert-demo, and held-out evaluation seed blocks.
• Random, masked-random, and greedy baseline policies.
• Exact solver utilities for sanity checks and solved examples.
• Queue provenance manifests and audits to check evaluation leakage.
• Static plots, GIF rollout rendering, and an interactive local dashboard.

Project Structure

The code is organized so puzzle mechanics live in one place:

• brainblock/core: board constants, pieces, action encoding, legality checks, state transitions, and feasibility checks.
• brainblock/mdp: Gymnasium environment, observations, and reward functions.
• brainblock/policies: random, masked-random, and greedy baseline policies.
• brainblock/agents/dqn: PyTorch Q-network, replay buffer, masked Double DQN update logic, and checkpoints.
• brainblock/solvers: exact backtracking solver and expert rollout helpers for sanity checks and examples.
• brainblock/experiment: training, evaluation, rollout, config, and metrics code.
• brainblock/visualization: ANSI/GIF rendering and learning-curve plots.
• brainblock/io: artifact path and manifest helpers.
• scripts: thin CLI entrypoints.

Action masking is used for DQN action selection and target computation, but the environment still exposes the full discrete action space.

Installation

The project uses Python 3.11+ and uv for dependency management:

uv sync --locked

The lockfile uses the CPU PyTorch index on Linux to avoid downloading CUDA wheels in CPU-only environments.

Install development tools:

uv sync --locked --extra dev

Development Checks

uv run ruff check brainblock scripts report/generate_figures.py
uv run ruff format --check brainblock scripts report/generate_figures.py
uv run basedpyright

Usage

Train

Run one reward mode and seed:

uv run python scripts/train.py --config configs/feasible.yaml --run-id dev --seed 0 --episodes 300

Useful shorter smoke run:

uv run python scripts/train.py \
  --config configs/feasible.yaml \
  --run-id smoke \
  --seed 0 \
  --episodes 80 \
  --expert-bootstrap-episodes 100 \
  --behavior-clone-steps 800 \
  --expert-offline-updates 300

Evaluate

Evaluate a trained DQN checkpoint:

uv run python scripts/evaluate.py --policy dqn --run-id dev --reward-mode feasible --seed 0 --episodes 50

By default, evaluation loads the fixed-budget final checkpoint, checkpoint.pt, and uses the checkpoint's configured action-mask mode. Pass --checkpoint-kind best to evaluate best_checkpoint.pt, which is selected by periodic validation; best_train_checkpoint.pt is kept only as a diagnostic best-training-return checkpoint. Evaluation episode seeds are offset from training seeds by default and evaluated in non-overlapping seed blocks. The evaluator skips duplicate queues and queues seen in training, validation, or expert demonstrations; pass --eval-seed-offset 0 only for debugging.

Evaluate all baseline policies:

uv run python scripts/evaluate.py --policy all-baselines --run-id dev --reward-mode feasible --seed 0 --episodes 50

Each DQN evaluation writes agent_solutions.json when solved rollouts are found. To collect solved agent rollouts across reward modes and seeds into the run-level solutions directory:

uv run python scripts/collect_agent_solutions.py --run-id dev --max-solutions 50

Rollout Demo

Render a solved rollout from the learned neural Q-network with the configured action mask:

uv run python scripts/demo.py \
  --run-id dev \
  --reward-mode feasible \
  --from-agent-solutions \
  --require-solved \
  --gif

This command infers the matching final checkpoint from solutions/agent_solutions.json, clears any expert-action table after loading, replays the stored held-out evaluation queue, and fails if the learned policy does not solve it. Pass --checkpoint-kind best only when intentionally replaying a validation-selected checkpoint. The optional --expert-policy-probability flag is rejected for this replay path.

Dashboard

Serve an interactive dashboard for inspecting a run:

uv run python scripts/dashboard.py --run-id main

Open http://127.0.0.1:8765. The dashboard lists solved held-out example queues, lets you type or shuffle inventory-valid queues, compares DQN against random, masked-random, and greedy policies, and renders each rollout frame by frame. DQN runs load the selected checkpoint and clear expert actions before inference.

Plots

uv run python scripts/plot_results.py --run-id dev

This writes:

• results/runs/dev/plots/learning_curves.png: one line per reward/seed run.
• results/runs/dev/plots/learning_curves_by_reward.png: one line per reward mode, averaging seeds with a standard-error band.
• results/runs/dev/plots/evaluation_by_reward.png: aggregate DQN evaluation comparison by reward mode.

The learning-curve figures include:

• total reward vs. episode
• total covered area vs. episode
• episode length vs. episode
• invalid-action rate vs. episode

Exact Solutions

Generate at least five solved rollouts with the exact solver:

uv run python scripts/generate_solutions.py --run-id dev --count 5 --seed 0

Full Pipeline

The full helper runs sparse, dense, and feasible DQN experiments, evaluates DQN and baselines after all training artifacts exist, generates solutions, plots curves, audits queue separation, and renders rollout artifacts:

RUN_ID=main SEEDS="0 1 2 3 4" EPISODES=20000 SEED_STRIDE=25000 EVAL_EPISODES=50 CHECKPOINT_KIND=final scripts/full_project_run.sh

For a quick representative run:

RUN_ID=smoke SEEDS="0" EPISODES=20 EVAL_EPISODES=5 REWARD_MODES="dense feasible" \
EXPERT_BOOTSTRAP_EPISODES=5 BEHAVIOR_CLONE_STEPS=50 EXPERT_OFFLINE_UPDATES=20 RUN_DEMO=0 scripts/full_project_run.sh

The dashboard is intentionally not part of the full helper because it is a local inspection UI rather than a batch experiment step. Run uv run python scripts/dashboard.py --run-id main separately when you want to inspect a completed run.

No-Expert Ablation

Run the no-expert feasible ablation:

for seed in 0 1 2 3 4; do
  uv run python scripts/train.py --config configs/no_expert_feasible.yaml --run-id no_expert_ablation --seed "$seed"
done
for seed in 0 1 2 3 4; do
  uv run python scripts/evaluate.py \
    --policy dqn \
    --run-id no_expert_ablation \
    --reward-mode feasible \
    --seed "$seed" \
    --checkpoint-kind final \
    --exclude-run-id main \
    --exclude-run-id no_expert_ablation
done
uv run python scripts/write_queue_manifest.py --run-id no_expert_ablation
uv run python scripts/collect_agent_solutions.py --run-id no_expert_ablation

Queue Audit

Audit generated queue provenance:

uv run python scripts/audit_queues.py --run-id main --run-id no_expert_ablation --fail-on-issue

Save the final audit artifact:

uv run python scripts/audit_queues.py \
  --run-id main \
  --run-id no_expert_ablation \
  --reference-run-id main \
  --output results/runs/queue_audit.json \
  --fail-on-issue

Report

Regenerate report figures and compile the PDF:

uv run python scripts/plot_results.py --run-id main
uv run python report/generate_figures.py
cd report
tex-fmt -n report.tex
latexmk -pdf report.tex
latexmk -c report.tex

If tex-fmt is unavailable, skip that line and run latexmk -pdf -bibtex report.tex followed by latexmk -c report.tex.

Learning-curve regeneration requires the training metrics.csv logs. Queue auditing does not require those raw logs: queue_manifest.json stores a compact source-queue record, and the audit script can also reconstruct training queues directly from saved configs.

Outputs

Run artifacts are written under results/runs/:

results/runs/
  queue_audit.json
  <run_id>/
    manifest.json
    queue_manifest.json
    training/<reward_mode>/seed_<seed>/
      config.json
      metrics.csv
      summary.json
      checkpoint.pt
      best_checkpoint.pt
      best_train_checkpoint.pt  # optional diagnostic checkpoint on fresh runs
      solved_rollouts.json
      expert_solutions.json
    evaluation/
      dqn/<reward_mode>/seed_<seed>/
        episodes.csv
        summary.json
        rollouts.json
        agent_solutions.json
      baselines/<policy_name>/<reward_mode>/seed_<seed>/
    solutions/
      solver_solutions.json
      agent_solutions.json
    plots/learning_curves.png
    demo/rollout.json
    demo/rollout.txt
    demo/rollout.gif

Re-running a command with the same --run-id overwrites files inside that run directory rather than creating timestamped folders.

Reproducibility Notes

Training and evaluation scripts accept explicit seeds. Python, NumPy, PyTorch, and the Gymnasium environment are seeded before agent construction and episode generation. Training, validation, expert-demonstration, and final-evaluation queues use separate seed blocks. Final evaluation also checks queue identity and skips any queue already seen in training, validation, or demonstration data. Configs compare three reward functions: sparse, dense, and feasible. The feasible reward adds simple connected-empty-component shaping that penalizes empty regions whose area is not divisible by four.

The main DQN pipeline can use a feasibility-safe action mask that filters legal actions which immediately create impossible empty components. It also warm-starts from exact-solver rollouts with behavior cloning and offline Double DQN updates before online fine-tuning. For ablation, configs/no_expert_feasible.yaml disables expert rollouts, behavior cloning, and expert replay.

The 20,000-episode reference artifacts were generated from the code revision recorded in the run manifests at commit 3818029. Later commits update packaging notes, report text, bibliography metadata, and reproducibility documentation without changing the reported evaluation artifacts.

If dependency resolution is changed later, keep the Linux PyTorch source pinned to the CPU index so clean CPU environments do not download CUDA wheels unnecessarily.

No RL framework is used. The DQN algorithm, replay buffer, target network update, action masking, gradient clipping, and checkpointing are implemented directly with PyTorch.