Add sparse maze option

build.sh

@@ -54,7 +54,15 @@ fi
 PLATFORM="$(uname -s)"
 if [ "$PLATFORM" = "Linux" ]; then
     RAYLIB_NAME='raylib-5.5_linux_amd64'
-    OMP_LIB=-lomp5
+    if ldconfig -p 2>/dev/null | grep -q 'libomp5\.so'; then
+        OMP_LIB=-lomp5
+    elif ldconfig -p 2>/dev/null | grep -q 'libomp\.so\.5'; then
+        OMP_LIB=-l:libomp.so.5
+    elif ldconfig -p 2>/dev/null | grep -q 'libomp\.so'; then
+        OMP_LIB=-lomp
+    else
+        OMP_LIB=-lgomp
+    fi
     SANITIZE_FLAGS=(-fsanitize=address,undefined,bounds,pointer-overflow,leak -fno-omit-frame-pointer)
     STANDALONE_LDFLAGS=(-lGL)
     SHARED_LDFLAGS=(-Bsymbolic-functions -Wl,--gc-sections)
@@ -172,6 +180,7 @@ fi
 CUDA_HOME=${CUDA_HOME:-${CUDA_PATH:-$(dirname "$(dirname "$(which nvcc)")")}}
 CUDNN_IFLAG=""
 CUDNN_LFLAG=""
+CUDNN_LIB="-lcudnn"
 for dir in /usr/local/cuda/include /usr/include; do
     if [ -f "$dir/cudnn.h" ]; then
         CUDNN_IFLAG="-I$dir"
@@ -190,6 +199,9 @@ fi
 if [ -z "$CUDNN_LFLAG" ]; then
     CUDNN_LFLAG=$(python -c "import nvidia.cudnn, os; print('-L' + os.path.join(nvidia.cudnn.__path__[0], 'lib'))" 2>/dev/null || echo "")
 fi
+if [ -n "$CUDNN_LFLAG" ] && [ ! -f "${CUDNN_LFLAG#-L}/libcudnn.so" ] && [ -f "${CUDNN_LFLAG#-L}/libcudnn.so.9" ]; then
+    CUDNN_LIB="-l:libcudnn.so.9"
+fi
 
 # NCCL include/lib fallback (mirrors the cuDNN fallback above).
 # Needed when NCCL is provided by the nvidia-nccl-cu12 wheel in the active venv.
@@ -218,7 +230,11 @@ done
 export CCACHE_DIR="${CCACHE_DIR:-$HOME/.ccache}"
 export CCACHE_BASEDIR="$(pwd)"
 export CCACHE_COMPILERCHECK=content
-NVCC="ccache $CUDA_HOME/bin/nvcc"
+if command -v ccache >/dev/null; then
+    NVCC="ccache $CUDA_HOME/bin/nvcc"
+else
+    NVCC="$CUDA_HOME/bin/nvcc"
+fi
 CC="${CC:-$(command -v ccache >/dev/null && echo 'ccache clang' || echo 'clang')}"
 ARCH=${NVCC_ARCH:-native}
 
@@ -273,7 +289,7 @@ if [ -z "$MODE" ]; then
         build/bindings.o "$RAYLIB_A"
         -L$CUDA_HOME/lib64 $CUDNN_LFLAG $NCCL_LFLAG
         "${WHEEL_RPATH_FLAGS[@]}"
-        -lcudart -lnccl -lnvidia-ml -lcublas -lcusolver -lcurand -lcudnn
+        -lcudart -lnccl -lnvidia-ml -lcublas -lcusolver -lcurand $CUDNN_LIB
         $OMP_LIB $LINK_OPT
         "${SHARED_LDFLAGS[@]}"
         -o "$OUTPUT"
@@ -314,7 +330,7 @@ elif [ "$MODE" = "profile" ]; then
         -Xcompiler=-fopenmp \
         tests/profile_kernels.cu vendor/ini.c \
         "$RAYLIB_A" \
-        -lnccl -lnvidia-ml -lcublas -lcurand -lcudnn \
+        -lnccl -lnvidia-ml -lcublas -lcurand $CUDNN_LIB \
         -lGL -lm -lpthread $OMP_LIB \
         -o profile
     echo "Built: ./profile"

config/maze.ini

@@ -10,6 +10,8 @@ num_threads = 0
 [env]
 num_maps = 8192
 map_size = 35
+sparse = 0
+wall_prob = 0.15
 
 [policy]
 hidden_size = 1024

ocean/maze/binding.c

@@ -20,6 +20,8 @@ Env* my_vec_init(int* num_envs_out, int* buffer_env_starts, int* buffer_env_coun
     int max_size = MAX_SIZE;
     int num_maps = (int)dict_get(env_kwargs, "num_maps")->value;
     int map_size = (int)dict_get(env_kwargs, "map_size")->value;
+    int sparse = (int)dict_get(env_kwargs, "sparse")->value;
+    float wall_prob = (float)dict_get(env_kwargs, "wall_prob")->value;
 
     if (max_size <= 5) {
         *num_envs_out = 0;
@@ -44,8 +46,12 @@ Env* my_vec_init(int* num_envs_out, int* buffer_env_starts, int* buffer_env_coun
         level->width = sz;
         level->height = sz;
 
-        float difficulty = (float)rand_r(&map_rng) / (float)(RAND_MAX);
-        create_maze_level(level, difficulty, i);
+        if (sparse) {
+            create_sparse_maze_level(level, wall_prob, i);
+        } else {
+            float difficulty = (float)rand_r(&map_rng) / (float)(RAND_MAX);
+            create_maze_level(level, difficulty, i);
+        }
     }
 
     // Allocate all environments

ocean/maze/maze.h

@@ -443,3 +443,31 @@ void create_maze_level(State* s, float difficulty, int seed) {
     int goal_adr = maze_offset(s->height - 2, s->width - 2);
     s->maze[goal_adr] = GOAL;
 }
+
+void create_sparse_maze_level(State* s, float wall_prob, int seed) {
+    unsigned int rng = seed;
+    memset(s->maze, EMPTY, MAX_SIZE*s->height);
+
+    for (int r = 1; r < s->height - 1; r++) {
+        for (int c = 1; c < s->width - 1; c++) {
+            float u = (float)rand_r(&rng) / (float)RAND_MAX;
+            if (u < wall_prob) {
+                s->maze[maze_offset(r, c)] = WALL;
+            }
+        }
+    }
+
+    make_border(s);
+
+    int goal_y = s->height - 2;
+    int goal_x = s->width - 2;
+    for (int c = 1; c <= goal_x; c++) {
+        s->maze[maze_offset(1, c)] = EMPTY;
+    }
+    for (int r = 1; r <= goal_y; r++) {
+        s->maze[maze_offset(r, goal_x)] = EMPTY;
+    }
+
+    spawn_agent(s, 0, 1, 1);
+    s->maze[maze_offset(goal_y, goal_x)] = GOAL;
+}