Skip to content

Add Craftax-Classic env to Ocean#531

Merged
jsuarez5341 merged 2 commits intoPufferAI:4.0from
Infatoshi:craftax-ocean
Apr 18, 2026
Merged

Add Craftax-Classic env to Ocean#531
jsuarez5341 merged 2 commits intoPufferAI:4.0from
Infatoshi:craftax-ocean

Conversation

@Infatoshi
Copy link
Copy Markdown
Contributor

Single-header port of Craftax-Classic (Matthews et al., ICML 2024).
Follows the ocean/drmario template: craftax.h + binding.c +
config/ocean/craftax.ini.

Observation: 1345 float32 (7x9 tiles x 21 channels one-hot + 22 scalar) — matches the JAX Craftax-Classic-Symbolic-v1 layout exactly.
Action: 1 discrete in [0, 17) (noop, 4 moves, do, sleep, 4 place, 3 make-pick, 3 make-sword).
Log: perf = achievements/22, score, episode_return, episode_length, plus one counter per achievement.

Training sanity

RTX PRO 6000 Blackwell, ~30 seconds of training, default config from this PR:

metric value
SPS 2.5M @ 99% GPU utilization
perf 0.69 (15/22 achievements unlocked)
episode_return 14.1
avg episode length ~312 steps

Agent cleanly learns wood → table → wood_pick → wood_sword → plant, progresses to drink / collect_stone / defeat_zombie. No NaN / divergence.

Implementation notes

  • AVX-512 Perlin noise worldgen is enabled via a function-level __attribute__((target(...))), so the env itself requires no build.sh flag changes.
  • Per-tile mob occupancy bitmaps give O(1) `has_mob_at` / collision queries.
  • All game-logic helpers already operated on a single env state in the upstream craftax.c project, so no per-env vs batched rework was needed.

Commits

  1. `ocean/craftax: Craftax-Classic env` — three files, additive, no core changes.
  2. `build.sh: NCCL include/lib fallback via nvidia-nccl-cu12 wheel` — independent fix that mirrors the existing cuDNN fallback. Without it, builds fail on systems where NCCL is provided by the `nvidia-nccl-cu12` wheel in the active venv rather than a system install. Feel free to drop / split into its own PR if preferred.

Future work

Happy to follow up in separate PRs:

  • Custom encoder that takes a compact uint8 observation and expands on GPU — Craftax's one-hot-per-tile structure means the dense Linear on 1345-float obs is mostly zero reads. Compact obs would cut PCIe traffic ~37x. Needs a way to plug a custom encoder / obs dtype per-env.
  • `precision = bf16` flag in `[train]` for tensor-core-friendly matmul on Blackwell / Hopper. Our 3.0-era benches showed 1.7x on small-model minibatch microbenches; would want to re-measure inside 4.0's tuned CUDA trainer before committing to it.

Both are noted here rather than shipped so the env PR stays minimal.

Provenance

Derived from https://github.com/Infatoshi/craftax.c (standalone C/AVX-512 port, 47.8M SPS env-only bench). The Ocean port strips that project's batched-OpenMP harness since `src/vecenv.h` already OMPs across env instances.

Single-header per-env port of Craftax-Classic (Matthews et al. ICML 2024),
following the ocean/drmario template.

  ocean/craftax/craftax.h     struct + c_init/reset/step/close/render,
                              AVX-512 Perlin worldgen, per-tile mob
                              occupancy bitmaps
  ocean/craftax/binding.c     OBS_SIZE=1345, NUM_ATNS=1, ACT_SIZES={17},
                              my_init, my_log
  config/ocean/craftax.ini    minimal (matches post-drmario template)

Observation: 1345 float32 (7x9 tiles x 21 channels + 22 scalar), matches
the JAX Craftax-Classic-Symbolic-v1 layout exactly.
Action: 1 discrete in [0, 17).
Log: perf (= achievements/22), score, episode_return, episode_length,
plus one counter per achievement.

Training sanity (RTX PRO 6000 Blackwell, 30s, default config):
  SPS:              2.5M at 99% GPU utilization
  perf:             0.69  (15/22 achievements unlocked)
  episode_return:   14.1, avg episode length 312
  Agent cleanly learns wood -> table -> wood_pick -> wood_sword -> plant,
  progresses to drink, collect_stone, defeat_zombie. No NaN/divergence.

AVX-512 in generate_world is enabled via function-level
__attribute__((target(...))), so the env itself does not require any
build.sh changes.

Derived from github.com/Infatoshi/craftax.c (standalone C/AVX-512 port,
47.8M SPS env-only bench). The Ocean port strips that project's batched
OpenMP harness since vecenv.h already OMPs across env instances.
Mirror the existing cuDNN fallback block for NCCL. When NCCL is provided
by the nvidia-nccl-cu12 wheel in the active venv rather than a system
install, the build otherwise fails with

    src/pufferlib.cu:5:10: fatal error: nccl.h: No such file or directory

Same search pattern as cuDNN: check common system paths first, then
fall back to python -c "import nvidia.nccl".
@jsuarez5341 jsuarez5341 merged commit 919609b into PufferAI:4.0 Apr 18, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants