Craftax GPU-resident env behind opt-in vecenv device hooks

[Infatoshi]

What this is

Craftax as PufferLib's first GPU-resident env, behind an opt-in vecenv hook. Companion PR #615 carries the env-agnostic trainer optimizations that were split out of this one; each stands alone. Everything here is additive: the default build is byte-identical in behavior (hash-verified), the device env is behind GPU_ENV=1 ./build.sh craftax, compact obs behind CRAFTAX_COMPACT_OBS.

All numbers: stock 4.0 trainer recipe as-is (single pass, 32 minibatches of 32768, bf16, Muon, V-trace + prio replay, cudagraphs), 8192 envs, MinGRU h256 x 3 layers, horizon 128, reset_pool_size=1024, RTX PRO 6000 Blackwell — measured on exactly this branch:

config	rollout SPS	train SPS
CPU env (16 buffers / 16 threads)	4.78M	3.46M
GPU env, this PR	24.3M	8.12M
GPU env + compact obs	25.3M	9.01M

(With #615 applied on top: 8.73M / 9.85M.)

Design: MY_GPU_ENV hooks (src/vecenv.h)

An env may define five hooks (my_gpu_init/reset/step/sync_logs/close). If it does, the threadmanager's per-step D2H-actions -> memset -> OpenMP c_step -> H2D block is replaced by a single async device call on the buffer stream — zero host syncs per step. Requires num_buffers=1. Host envs still build, so the CPU path stays available in the same process (that's how parity is tested below). Envs that don't define the hooks are completely untouched. This is the same direction as the (currently unconsumed) MY_VEC_STEP macros in the 5.0 craftax binding — happy to converge on whatever shape 5.0 settles on.

The env

ocean/craftax/craftax.cu is a vendored env-only snapshot of craftax.cu (single-file CUDA port of full Craftax, -DCRAFTAX_CU_LIB gates its main) — the device sibling of craftax.c/craftax.h. The game logic and step/reset/encode kernels are verbatim; the standalone repo's own trainer/bench harness (~4.6k lines PufferLib never calls) is stripped by a mechanical reachability script that lives upstream, so the snapshot is 9.4k lines and regenerable in one command. The trimmed build was gated the same way as everything else: byte-identical trajectory hashes (512x300, 4096x200, pool-1024 4096x200) and a bitwise-identical 30-iteration learning trace against the untrimmed build. binding.cu wraps it behind the hooks (device sibling of binding.c), including device reset pools matching the host reset_pool_size semantics (pre-generated worlds, reset = state copy; index selection matches the host done-reset path).

Correctness

• CPU-vs-GPU bitwise parity: same _C.so, same seeding (rng=i, seed=seed_offset+i, zero glue), same action stream — obs/rewards/terminals bitwise-identical every step, verified 256x500 and 4096x300 (tests/craftax_gpu_parity.py). Caveat stated plainly: bitwise holds when the C binding is compiled with the reference flags (gcc -O3 -ffast-math); clang -O2 differs by 1-ULP float codegen in ~5 scalar obs floats — the same compiler-relative variation that already exists between clang/gcc builds of the C env itself.
• Every behavior-invisible optimization (pool resets, copy slimming) was gated on byte-identical per-step trajectory hashes before merging.

Compact observations (opt-in)

Same 843-value packed-ID layout the env already emits — map cells stored as uint8 tokens (996 bytes vs 3372) with a fixed lookup expansion in the compiled encoder (CraftaxCompactEncoder on the Python side). Besides throughput, it learns better on this recipe: 5 seeds x 5B steps each, compact final score 19.0 +/- 0.35 (all runs monotonic) vs float 11.6 (range 6.6-17.6, with mid-training policy collapse in 3 of 5 seeds). Default stays float32.

Limitations

• GPU env requires num_buffers=1 (nothing to overlap when the env steps on-device).
• Tested on sm_120 and sm_86. Env counts: multiples of 32 (warp record kernel); benchmarked at powers of two.
• 5.0's curriculum captures/restores host State; a device-resident env there needs a small state bridge (or curriculum off) — noted for the eventual port, not relevant to 4.0.

🤖 Generated with Claude Code