Skip to content

Add Craftax moonshot environment#539

Closed
Infatoshi wants to merge 2 commits into
PufferAI:4.0from
Infatoshi:craftax-moonshot-pr
Closed

Add Craftax moonshot environment#539
Infatoshi wants to merge 2 commits into
PufferAI:4.0from
Infatoshi:craftax-moonshot-pr

Conversation

@Infatoshi
Copy link
Copy Markdown
Contributor

Summary

Adds a side-by-side craftax_moonshot implementation for the 4.0 branch. This keeps the existing Craftax implementation intact while adding a CPU-oriented version with a smaller packed float observation path, reset-state caching, tiled stepping, and a per-buffer VecEnv rollout hook used only by environments that opt in.

The src/vecenv.h change is intentionally narrow: default environments keep the existing c_step rollout loop, while craftax_moonshot defines MY_VEC_STEP_RANGE to avoid the scalar fallback during Puffer rollouts.

Measured throughput

On a Ryzen 9 9950X3D + RTX PRO 6000 Blackwell workstation, pinned to the V-Cache cores:

env OMP_NUM_THREADS=1 taskset -c 0-7,16-23 \
  uv run puffer train craftax_moonshot \
  --train.total-timesteps 200000000 \
  --vec.total-agents 16384 \
  --vec.num-buffers 16 \
  --vec.num-threads 16 \
  --policy.hidden-size 32 \
  --policy.num-layers 1 \
  --policy.expansion-factor 1 \
  --train.horizon 128 \
  --train.minibatch-size 32768

Short-run dashboard measurement: about 4.0M training SPS with env around 47% of wall time.

Raw random-action env-only measurement on 4096 agents: about 16.8M env SPS.

Verification

uv run --with pybind11 --with rich_argparse --with nvidia-cudnn-cu12 --with nvidia-nccl-cu12 ./build.sh craftax_moonshot
uv run --with pytest pytest tests/craftax_diff.py -q
uv run --with ruff ruff check tests/craftax_diff.py --fix

Results:

Built: pufferlib/_C.cpython-313-x86_64-linux-gnu.so
2 passed in 15.07s
All checks passed!

Repo-wide pytest/ruff are currently blocked by pre-existing upstream dependency/lint issues unrelated to this PR.

@Infatoshi
Copy link
Copy Markdown
Contributor Author

Closing this version. Reworking the patch to update the existing craftax environment directly instead of adding craftax_moonshot side-by-side.

@Infatoshi Infatoshi closed this Apr 24, 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.

1 participant