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

build.sh

@@ -140,6 +140,13 @@ fi
 
 OUTPUT_NAME=${OUTPUT_NAME:-$ENV}
 
+# Craftax's compact observation ABI must be selected consistently in the host
+# binding, optional device env, and compiled trainer. Leave it opt-in so the
+# default float-observation build remains unchanged.
+if [ "${COMPACT_OBS:-0}" = "1" ]; then
+    EXTRA_CFLAGS="$EXTRA_CFLAGS -DCRAFTAX_COMPACT_OBS"
+fi
+
 # Standalone environment build
 # -mavx2 enables AVX2 intrinsics (__m256, _mm256_*) which drive.h and
 # src/bf16.h use directly. x86_64 only — strip if porting to ARM/Apple Silicon.
@@ -258,6 +265,33 @@ if [ ! -f "$BINDING_SRC" ]; then
     exit 1
 fi
 
+# Optional device-resident env: GPU_ENV=1 compiles $SRC_DIR/binding.cu
+# (a CUDA implementation of the my_gpu_* hooks in src/vecenv.h) into the env
+# static lib and defines PUFFER_GPU_ENV for the binding so the vec steps the
+# env entirely on the GPU. CUDA build mode only.
+GPU_ENV_SRC="$SRC_DIR/binding.cu"
+GPU_ENV_OBJ=""
+if [ "${GPU_ENV:-0}" = "1" ]; then
+    if [ ! -f "$GPU_ENV_SRC" ]; then
+        echo "Error: GPU_ENV=1 but $GPU_ENV_SRC not found"
+        exit 1
+    fi
+    if [ -n "$MODE" ]; then
+        echo "Error: GPU_ENV=1 requires the default CUDA build mode"
+        exit 1
+    fi
+    EXTRA_CFLAGS="$EXTRA_CFLAGS -DPUFFER_GPU_ENV"
+    GPU_ENV_OBJ="build/libgpu_${ENV}.o"
+    echo "Compiling device env $GPU_ENV_SRC..."
+    $NVCC -c -arch=$ARCH -O3 -Xcompiler -fPIC \
+        --expt-relaxed-constexpr -fmad=false \
+        -std=c++17 \
+        -I. -Isrc -I$SRC_DIR \
+        -I$CUDA_HOME/include \
+        $PRECISION $EXTRA_CFLAGS \
+        "$GPU_ENV_SRC" -o "$GPU_ENV_OBJ"
+fi
+
 echo "Compiling static library for $ENV..."
 ${CC:-clang} -c "${CLANG_OPT[@]}" $EXTRA_CFLAGS \
     -I. -Isrc -I$SRC_DIR -Ivendor \
@@ -267,10 +301,17 @@ ${CC:-clang} -c "${CLANG_OPT[@]}" $EXTRA_CFLAGS \
     -fno-semantic-interposition -fvisibility=hidden \
     -fPIC -fopenmp \
     "$BINDING_SRC" -o "$STATIC_OBJ"
-ar rcs "$STATIC_LIB" "$STATIC_OBJ"
+ar rcs "$STATIC_LIB" "$STATIC_OBJ" $GPU_ENV_OBJ
 
-# Brittle hack: have to extract the tensor type from the static lib to build trainer
-OBS_TENSOR_T=$(awk '/^#define OBS_TENSOR_T/{print $3}' "$BINDING_SRC")
+# Brittle hack: have to extract the tensor type from the static lib to build trainer.
+# Preprocess (with the same flags as the binding compile) so flag-gated defines
+# like -DCRAFTAX_COMPACT_OBS resolve to the tensor type actually compiled in.
+OBS_TENSOR_T=$(${CC:-clang} -E -dM $EXTRA_CFLAGS \
+    -I. -Isrc -I$SRC_DIR -Ivendor \
+    "${INCLUDES[@]}" \
+    -I./$RAYLIB_NAME/include -I$CUDA_HOME/include \
+    -DPLATFORM_DESKTOP \
+    "$BINDING_SRC" | awk '$2=="OBS_TENSOR_T"{print $3}')
 if [ -z "$OBS_TENSOR_T" ]; then
     echo "Error: Could not find OBS_TENSOR_T in $BINDING_SRC"
     exit 1
@@ -289,7 +330,7 @@ if [ -z "$MODE" ]; then
         -Xcompiler=-fopenmp \
         -DOBS_TENSOR_T=$OBS_TENSOR_T \
         -DENV_NAME=$ENV \
-        $PRECISION $NVCC_OPT \
+        $PRECISION $EXTRA_CFLAGS $NVCC_OPT \
         src/bindings.cu -o build/bindings.o
 
     LINK_CMD=(
@@ -316,7 +357,7 @@ elif [ "$MODE" = "cpu" ]; then
         -I$PYTHON_INCLUDE -I$PYBIND_INCLUDE \
         -DOBS_TENSOR_T=$OBS_TENSOR_T \
         -DENV_NAME=$ENV \
-        $PRECISION $LINK_OPT \
+        $PRECISION $EXTRA_CFLAGS $LINK_OPT \
         src/bindings_cpu.cpp -o build/bindings_cpu.o
     LINK_CMD=(
         ${CXX:-g++} -shared -fPIC -fopenmp
@@ -337,7 +378,7 @@ elif [ "$MODE" = "profile" ]; then
         -DOBS_TENSOR_T=$OBS_TENSOR_T \
         -DENV_NAME=$ENV \
         -Xcompiler=-DPLATFORM_DESKTOP \
-        $PRECISION \
+        $PRECISION $EXTRA_CFLAGS \
         -Xcompiler=-fopenmp \
         tests/profile_kernels.cu vendor/ini.c \
         "$STATIC_LIB" "$RAYLIB_A" \

config/craftax.ini

@@ -14,6 +14,14 @@ seed_offset = 0
 # ever seen per process. Set to 0 to disable (required for the
 # parity harness to maintain exact per-seed determinism).
 reset_pool_size = 1024
+# Lazy floor generation: resets generate only floor 0; floors 1-8 are
+# generated on first descent (bit-identical maps, deferred cost). Cheap
+# resets with unbounded world diversity; pair with reset_pool_size = 0.
+lazy_floors = 0
+
+# Compact byte observations (996 uint8 vs 843 float32, 3.4x less DMA):
+# rebuild with EXTRA_CFLAGS=-DCRAFTAX_COMPACT_OBS ./build.sh craftax
+# and set under [torch]: encoder = CraftaxCompactEncoder
 
 [train]
 total_timesteps = 200_000_000

ocean/craftax/binding.c

@@ -7,16 +7,46 @@
 #define OBS_SIZE CRAFTAX_OBS_SIZE
 #define NUM_ATNS 1
 #define ACT_SIZES {CRAFTAX_NUM_ACTIONS}
+#ifdef CRAFTAX_COMPACT_OBS
+#define OBS_TENSOR_T ByteTensor
+#else
 #define OBS_TENSOR_T FloatTensor
+#endif
 
 #define CRAFTAX_VEC_TILE_SIZE 128
 #define MY_VEC_INIT
 #define MY_VEC_CLOSE
 #define MY_VEC_STEP craftax_vec_step
 #define MY_VEC_STEP_RANGE craftax_vec_step_range
+#ifdef PUFFER_GPU_ENV
+// Device-resident env (binding.cu): the vec steps entirely on the GPU.
+#define MY_GPU_ENV 1
+#endif
 #define Env Craftax
 #include "vecenv.h"
 
+#ifdef PUFFER_GPU_ENV
+// Implemented in binding.cu: drains per-env device logs (same Log layout
+// as the CPU env) and zeroes them device-side.
+void craftax_gpu_collect_logs(void* logs_out, int num_envs);
+
+void my_gpu_sync_logs(StaticVec* vec) {
+    Craftax* envs = (Craftax*)vec->envs;
+    int n = vec->size;
+    Log* logs = (Log*)malloc((size_t)n * sizeof(Log));
+    craftax_gpu_collect_logs(logs, n);
+    int num_keys = (int)(sizeof(Log) / sizeof(float));
+    for (int i = 0; i < n; i++) {
+        float* dst = (float*)&envs[i].log;
+        float* src = (float*)&logs[i];
+        for (int j = 0; j < num_keys; j++) {
+            dst[j] += src[j];
+        }
+    }
+    free(logs);
+}
+#endif
+
 // Tiled vector step: process agents in tiles that fit comfortably in cache.
 // Each thread processes a contiguous block of lightweight env handles while
 // the heavier CraftaxState storage lives in a separate arena.
@@ -130,6 +160,26 @@ void my_init(Env* env, Dict* kwargs) {
     if (pool_item != NULL) reset_pool_size = (int)pool_item->value;
     craftax_set_reset_pool_size(reset_pool_size);
 
+    // Lazy floor generation: reset generates only floor 0, floors 1..8 are
+    // generated on first descent (bit-identical maps, deferred cost).
+    int lazy_floors = 0;
+    DictItem* lazy_item = dict_get_unsafe(kwargs, "lazy_floors");
+    if (lazy_item != NULL) lazy_floors = (int)lazy_item->value;
+    craftax_set_lazy_floors(lazy_floors);
+
+    // Pipelined world pool: producer threads pre-generate fresh worlds so
+    // episode resets become a memcpy. Every world is unique (unlike the
+    // static reset pool) but reset outcomes depend on thread scheduling.
+    int wp_producers = 0;
+    int wp_capacity = 1024;
+    DictItem* wp_item = dict_get_unsafe(kwargs, "world_pool_producers");
+    if (wp_item != NULL) wp_producers = (int)wp_item->value;
+    DictItem* wpc_item = dict_get_unsafe(kwargs, "world_pool_capacity");
+    if (wpc_item != NULL) wp_capacity = (int)wpc_item->value;
+    if (wp_producers > 0) {
+        craftax_world_pool_start(wp_capacity, wp_producers, env->seed);
+    }
+
     c_init(env);
 }