Skip to content

4.0#402

Open
jsuarez5341 wants to merge 628 commits into3.0from
4.0
Open

4.0#402
jsuarez5341 wants to merge 628 commits into3.0from
4.0

Conversation

@jsuarez5341
Copy link
Copy Markdown
Contributor

This PR will not be merged. We are targeting EoY and 4.0 will just become master. Key goals:

  • Sweeps for all envs, largest ever public dataset of RL experiments
  • Constellation
  • Major perf enhancements

TBD: cpp/barracuda, final constellation features, xlstm, advantage calc tweaks

jsuarez5341 and others added 30 commits April 5, 2026 12:41
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".
Add Craftax-Classic env to Ocean
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.

8 participants