hyperray 0.0.1__py3-none-any.whl

hyperdrone/_vendor/rl-tools/src/rendering/raytracing/yaw_prediction/yaw_prediction_viewer.cu

@@ -0,0 +1,443 @@
+#define RL_TOOLS_OPERATIONS_CPU_MUX_INCLUDE_CUDA
+#include <rl_tools/operations/cpu_mux.h>
+
+// Layer operations
+#include <rl_tools/nn/layers/dense/operations_generic.h>
+#include <rl_tools/nn/layers/conv2d/operations_generic.h>
+#include <rl_tools/nn/layers/avg_pool2d/operations_generic.h>
+#include <rl_tools/nn/layers/dynamic_conv2d/operations_generic.h>
+#include <rl_tools/nn/layers/dynamic_conv2d/operations_cuda.h>
+
+// Model operations
+#include <rl_tools/nn_models/sequential/operations_generic.h>
+
+// CUDA-specific operations
+#include <rl_tools/nn/operations_cuda.h>
+#include <rl_tools/nn_models/operations_generic.h>
+
+// Tensor operations
+#include <rl_tools/containers/tensor/operations_generic.h>
+#include <rl_tools/containers/tensor/operations_cpu.h>
+
+// Persist (HDF5)
+#include <rl_tools/persist/backends/hdf5/operations_cpu.h>
+#include <rl_tools/nn/layers/dense/persist.h>
+#include <rl_tools/nn/layers/conv2d/persist.h>
+#include <rl_tools/nn/layers/avg_pool2d/persist.h>
+#include <rl_tools/nn/layers/dynamic_conv2d/persist.h>
+#include <rl_tools/nn_models/sequential/persist.h>
+#include <rl_tools/nn_models/parallel/persist.h>
+
+// Model definition
+#include "model.h"
+#include "model_operations.h"
+#include "model_persist.h"
+#include "model_forward_cuda.h"
+
+// Scene management
+#include "scene.h"
+
+#include <iostream>
+#include <vector>
+#include <cmath>
+#include <cstring>
+#include <sstream>
+#include <iomanip>
+#include <algorithm>
+#include <filesystem>
+#include <cuda_runtime.h>
+#include <cstdio>
+
+namespace rlt = rl_tools;
+namespace fs = std::filesystem;
+
+// ---- Configuration ----
+using T = float;
+using DEVICE_CPU = rlt::devices::DEVICE_FACTORY<>;
+using DEVICE_CUDA = rlt::devices::DEVICE_FACTORY_CUDA<>;
+using TI = DEVICE_CPU::index_t;
+using TI_CUDA = DEVICE_CUDA::index_t;
+
+namespace yp = rlt::rendering::raytracing::yaw_prediction;
+static constexpr TI_CUDA BATCH_SIZE = yp::SCENE_NUM_CAMERAS / 2; // 128
+static constexpr TI_CUDA NUM_CAMERAS = yp::SCENE_NUM_CAMERAS;    // 256
+static constexpr TI CAM_WIDTH = yp::SCENE_CAM_WIDTH;
+static constexpr TI CAM_HEIGHT = yp::SCENE_CAM_HEIGHT;
+
+static constexpr float PI = 3.14159265358979323846f;
+static constexpr float MAX_ANGLE = PI / 6.0f; // 30 degrees
+static constexpr float VIEWER_FOV = 37.81521447863434f; // Fixed FOV for viewer (degrees)
+static constexpr TI NUM_GRID_CAMERAS = 64;
+static constexpr TI GRID_DIM = 8; // 8x8 grid
+
+using TYPE_POLICY = rlt::numeric_types::Policy<float>;
+
+// Inference-only model
+using GPU_CAPABILITY = rlt::nn::capability::Forward<>;
+using GPU_MODEL = yp::MODEL<GPU_CAPABILITY, TYPE_POLICY, TI_CUDA, BATCH_SIZE, CAM_HEIGHT, CAM_WIDTH>;
+using CPU_CAPABILITY = rlt::nn::capability::Forward<>;
+using CPU_MODEL = yp::MODEL<CPU_CAPABILITY, TYPE_POLICY, TI, BATCH_SIZE, CAM_HEIGHT, CAM_WIDTH>;
+
+using GPU_INPUT_SHAPE = rlt::tensor::Shape<TI_CUDA, BATCH_SIZE, CAM_HEIGHT, CAM_WIDTH, 3>;
+using GPU_INPUT_SPEC = rlt::tensor::Specification<T, TI_CUDA, GPU_INPUT_SHAPE>;
+
+// Head output tensor: [BATCH_SIZE, 3]
+using GPU_OUTPUT_SHAPE = typename GPU_MODEL::OUTPUT_SHAPE;
+using GPU_OUTPUT_SPEC = rlt::tensor::Specification<T, TI_CUDA, GPU_OUTPUT_SHAPE>;
+
+// ---- Progress bead drawing ----
+void draw_bead(std::vector<uint8_t>& image, int image_width, int tile_x, int tile_y,
+               float angle, float max_angle, int bead_y, uint8_t r, uint8_t g, uint8_t b) {
+    float normalized = (angle + max_angle) / (2.0f * max_angle);
+    normalized = std::max(0.0f, std::min(1.0f, normalized));
+    int bead_x = 2 + static_cast<int>(normalized * 59.0f);
+    for (int dy = 0; dy < 3; dy++) {
+        for (int dx = 0; dx < 5; dx++) {
+            int px = tile_x + bead_x + dx;
+            int py = tile_y + bead_y + dy;
+            int idx = (py * image_width + px) * 3;
+            image[idx + 0] = r;
+            image[idx + 1] = g;
+            image[idx + 2] = b;
+        }
+    }
+}
+
+int main(int argc, char** argv) {
+    std::string scene_dir;
+    std::string single_scene;
+    const char* checkpoint_path = nullptr;
+    const char* output_path = "yaw_prediction_viewer.mp4";
+    int num_frames = 120;
+    TI num_scenes = 64;
+
+    for (int i = 1; i < argc; i++) {
+        std::string arg = argv[i];
+        if (arg == "--scene-dir" && i + 1 < argc) {
+            scene_dir = argv[++i];
+        } else if (arg == "--scene" && i + 1 < argc) {
+            single_scene = argv[++i];
+        } else if (arg == "--checkpoint" && i + 1 < argc) {
+            checkpoint_path = argv[++i];
+        } else if (arg == "--output" && i + 1 < argc) {
+            output_path = argv[++i];
+        } else if (arg == "--num-frames" && i + 1 < argc) {
+            num_frames = std::atoi(argv[++i]);
+        } else if (arg == "--num-scenes" && i + 1 < argc) {
+            num_scenes = std::atoi(argv[++i]);
+        } else {
+            std::cerr << "Usage: " << argv[0]
+                      << " [--scene-dir <dir> | --scene <path.glb>]"
+                      << " --checkpoint <path.h5>"
+                      << " [--output <path.mp4>] [--num-frames <N>]"
+                      << " [--num-scenes <N>]" << std::endl;
+            return 1;
+        }
+    }
+
+    if (!checkpoint_path) {
+        std::cerr << "Error: --checkpoint is required" << std::endl;
+        return 1;
+    }
+
+    // Collect scene paths
+    std::vector<std::string> all_scene_paths;
+    if (!scene_dir.empty()) {
+        for (auto& entry : fs::directory_iterator(scene_dir)) {
+            if (entry.path().extension() == ".glb") {
+                all_scene_paths.push_back(entry.path().string());
+            }
+        }
+        std::sort(all_scene_paths.begin(), all_scene_paths.end());
+        if (all_scene_paths.empty()) {
+            std::cerr << "No .glb files found in " << scene_dir << std::endl;
+            return 1;
+        }
+        if (num_scenes < all_scene_paths.size()) {
+            all_scene_paths.resize(num_scenes);
+        }
+    } else if (!single_scene.empty()) {
+        all_scene_paths.push_back(single_scene);
+    } else {
+        std::cerr << "Must specify --scene-dir or --scene" << std::endl;
+        return 1;
+    }
+
+    // Cap to 64 scenes (one per grid tile max)
+    if (all_scene_paths.size() > NUM_GRID_CAMERAS) {
+        all_scene_paths.resize(NUM_GRID_CAMERAS);
+    }
+
+    std::cout << "Yaw prediction viewer (multi-scene)" << std::endl;
+    std::cout << "  Scenes: " << all_scene_paths.size() << std::endl;
+    std::cout << "  Checkpoint: " << checkpoint_path << std::endl;
+    std::cout << "  Output: " << output_path << std::endl;
+    std::cout << "  Num frames: " << num_frames << std::endl;
+
+    // ---- Load scenes ----
+    struct LoadedScene {
+        yp::SceneHandle* handle = nullptr;
+        std::string path;
+    };
+    std::vector<LoadedScene> loaded_scenes(all_scene_paths.size());
+    for (TI s = 0; s < all_scene_paths.size(); s++) {
+        loaded_scenes[s].path = all_scene_paths[s];
+        std::cout << "Loading scene " << s << "/" << all_scene_paths.size() << ": "
+                  << fs::path(loaded_scenes[s].path).filename().string() << std::endl;
+        loaded_scenes[s].handle = yp::create_scene(loaded_scenes[s].path.c_str());
+        std::cout << "  Indoor states: " << yp::get_num_indoor_states(loaded_scenes[s].handle) << std::endl;
+    }
+
+    // Assign each grid tile to a scene (round-robin)
+    std::vector<TI> tile_scene_idx(NUM_GRID_CAMERAS);
+    for (TI i = 0; i < NUM_GRID_CAMERAS; i++) {
+        tile_scene_idx[i] = i % loaded_scenes.size();
+    }
+
+    // Group tiles by scene
+    std::vector<std::vector<TI>> scene_tiles(loaded_scenes.size());
+    for (TI i = 0; i < NUM_GRID_CAMERAS; i++) {
+        scene_tiles[tile_scene_idx[i]].push_back(i);
+    }
+
+    // ---- Device setup ----
+    DEVICE_CPU device_cpu;
+    DEVICE_CUDA device_cuda;
+    rlt::init(device_cpu);
+    rlt::init(device_cuda);
+
+    // ---- Load model ----
+    CPU_MODEL model_cpu;
+    rlt::malloc(device_cpu, model_cpu);
+    {
+        auto file = rl_tools::persist::backends::hdf5::File(checkpoint_path, rl_tools::persist::backends::hdf5::Mode::READ);
+        auto mg = rlt::get_group(device_cpu, file, "model");
+        bool success = rlt::load(device_cpu, model_cpu, mg);
+        if (!success) {
+            std::cerr << "Failed to load checkpoint" << std::endl;
+            return 1;
+        }
+    }
+    std::cout << "Checkpoint loaded" << std::endl;
+
+    GPU_MODEL model;
+    typename GPU_MODEL::template Buffer<true> model_buffer;
+    rlt::malloc(device_cuda, model);
+    rlt::malloc(device_cuda, model_buffer);
+    rlt::copy(device_cpu, device_cuda, model_cpu, model);
+
+    // ---- GPU tensors ----
+    rlt::Tensor<GPU_INPUT_SPEC> gpu_input_a, gpu_input_b;
+    rlt::malloc(device_cuda, gpu_input_a);
+    rlt::malloc(device_cuda, gpu_input_b);
+
+    rlt::Tensor<GPU_OUTPUT_SPEC> gpu_output;
+    rlt::malloc(device_cuda, gpu_output);
+
+    DEVICE_CUDA::SPEC::RANDOM::ENGINE<> rng_cuda;
+    rlt::malloc(device_cuda, rng_cuda);
+    rlt::init(device_cuda, rng_cuda, 42);
+
+    // ---- Sample base cameras per scene ----
+    std::vector<rlt::rendering::raytracing::Camera<float>> base_cameras(NUM_GRID_CAMERAS);
+    {
+        std::vector<rlt::rendering::raytracing::Camera<float>> setup_cameras(NUM_CAMERAS);
+        std::vector<float> setup_targets(BATCH_SIZE * 3);
+
+        for (TI s = 0; s < loaded_scenes.size(); s++) {
+            auto& tiles = scene_tiles[s];
+            if (tiles.empty()) continue;
+
+            yp::sample_camera_batch(loaded_scenes[s].handle, setup_cameras.data(),
+                                    setup_targets.data(),
+                                    BATCH_SIZE, 0.0f,
+                                    VIEWER_FOV, VIEWER_FOV);
+
+            for (TI t = 0; t < tiles.size(); t++) {
+                base_cameras[tiles[t]] = setup_cameras[t % BATCH_SIZE];
+            }
+        }
+    }
+
+    // Rotate a camera's view direction by delta_yaw around the Y axis
+    using YPCameraData = rlt::rendering::raytracing::Camera<float>;
+    auto rotate_camera_yaw = [](const YPCameraData& cam, float delta_yaw) -> YPCameraData {
+        YPCameraData rotated = cam;
+        float cy = std::cos(delta_yaw);
+        float sy = std::sin(delta_yaw);
+        auto rotate_y = [cy, sy](const float v[3], float out[3]) {
+            out[0] = cy * v[0] - sy * v[2];
+            out[1] = v[1];
+            out[2] = sy * v[0] + cy * v[2];
+        };
+        rotate_y(cam.dir_00, rotated.dir_00);
+        rotate_y(cam.dir_du, rotated.dir_du);
+        rotate_y(cam.dir_dv, rotated.dir_dv);
+        return rotated;
+    };
+
+    // ---- Open ffmpeg pipe ----
+    static constexpr TI IMAGE_WIDTH = GRID_DIM * CAM_WIDTH;
+    static constexpr TI IMAGE_HEIGHT = GRID_DIM * CAM_HEIGHT;
+    std::ostringstream ffmpeg_cmd;
+    ffmpeg_cmd
+        << "ffmpeg -y -f rawvideo -pixel_format rgb24 "
+        << "-video_size " << IMAGE_WIDTH << "x" << IMAGE_HEIGHT << " "
+        << "-framerate 30 -i - "
+        << "-an -c:v libx264 -pix_fmt yuv420p "
+        << output_path;
+    FILE* mp4_pipe = popen(ffmpeg_cmd.str().c_str(), "w");
+    if (!mp4_pipe) {
+        std::cerr << "Failed to start ffmpeg" << std::endl;
+        return 1;
+    }
+
+    // ---- Frame generation loop ----
+    static constexpr TI CAM_PIXELS = CAM_WIDTH * CAM_HEIGHT;
+    std::vector<uint32_t> host_fb(NUM_CAMERAS * CAM_PIXELS);
+    std::vector<uint8_t> output_image(IMAGE_WIDTH * IMAGE_HEIGHT * 3);
+    static constexpr TI OUTPUT_DIM = 3;
+    std::vector<float> pred_buf(BATCH_SIZE * OUTPUT_DIM);
+    std::vector<rlt::rendering::raytracing::Camera<float>> cameras(NUM_CAMERAS);
+
+    // Per-tile host framebuffer for swept views (read back per scene)
+    std::vector<std::vector<uint32_t>> tile_pixels(NUM_GRID_CAMERAS, std::vector<uint32_t>(CAM_PIXELS));
+
+    auto eval_mode = rlt::Mode<rlt::mode::Default<>>{};
+
+    for (int frame = 0; frame < num_frames; frame++) {
+        float t = (num_frames > 1) ? static_cast<float>(frame) / static_cast<float>(num_frames - 1) : 0.5f;
+        float delta_yaw = -MAX_ANGLE + t * 2.0f * MAX_ANGLE;
+
+        // ---- Render each scene and fill GPU input tensors ----
+        for (TI s = 0; s < loaded_scenes.size(); s++) {
+            auto& tiles = scene_tiles[s];
+            if (tiles.empty()) continue;
+
+            TI num_pairs = tiles.size();
+
+            for (TI t_idx = 0; t_idx < num_pairs; t_idx++) {
+                cameras[t_idx] = base_cameras[tiles[t_idx]];
+                cameras[num_pairs + t_idx] = rotate_camera_yaw(base_cameras[tiles[t_idx]], delta_yaw);
+            }
+            for (TI i = num_pairs * 2; i < NUM_CAMERAS; i++) {
+                cameras[i] = cameras[i % (num_pairs * 2)];
+            }
+
+            yp::render_batch<false>(loaded_scenes[s].handle, cameras.data());
+            uint32_t* device_fb = yp::get_framebuffer_device_ptr(loaded_scenes[s].handle);
+
+            for (TI t_idx = 0; t_idx < num_pairs; t_idx++) {
+                TI batch_pos = tiles[t_idx];
+                yp::cuda_kernels::rgba_to_activation_kernel<<<1, 256, 0, device_cuda.stream>>>(
+                    device_fb, gpu_input_a._data, CAM_WIDTH, CAM_HEIGHT, t_idx, batch_pos
+                );
+            }
+            for (TI t_idx = 0; t_idx < num_pairs; t_idx++) {
+                TI batch_pos = tiles[t_idx];
+                yp::cuda_kernels::rgba_to_activation_kernel<<<1, 256, 0, device_cuda.stream>>>(
+                    device_fb, gpu_input_b._data, CAM_WIDTH, CAM_HEIGHT, num_pairs + t_idx, batch_pos
+                );
+            }
+
+            cudaStreamSynchronize(device_cuda.stream);
+            for (TI t_idx = 0; t_idx < num_pairs; t_idx++) {
+                TI cam_idx = num_pairs + t_idx;
+                cudaMemcpy(tile_pixels[tiles[t_idx]].data(),
+                           device_fb + cam_idx * CAM_PIXELS,
+                           CAM_PIXELS * sizeof(uint32_t),
+                           cudaMemcpyDeviceToHost);
+            }
+        }
+
+        // Pad batch positions [64..127]
+        for (TI i = NUM_GRID_CAMERAS; i < BATCH_SIZE; i++) {
+            TI src_pos = i % NUM_GRID_CAMERAS;
+            cudaMemcpyAsync(
+                gpu_input_a._data + i * CAM_HEIGHT * CAM_WIDTH * 3,
+                gpu_input_a._data + src_pos * CAM_HEIGHT * CAM_WIDTH * 3,
+                CAM_HEIGHT * CAM_WIDTH * 3 * sizeof(float),
+                cudaMemcpyDeviceToDevice, device_cuda.stream
+            );
+            cudaMemcpyAsync(
+                gpu_input_b._data + i * CAM_HEIGHT * CAM_WIDTH * 3,
+                gpu_input_b._data + src_pos * CAM_HEIGHT * CAM_WIDTH * 3,
+                CAM_HEIGHT * CAM_WIDTH * 3 * sizeof(float),
+                cudaMemcpyDeviceToDevice, device_cuda.stream
+            );
+        }
+
+        // ---- Model forward pass ----
+        rlt::evaluate_cuda(device_cuda, model, gpu_input_a, gpu_input_b, gpu_output, model_buffer, rng_cuda, eval_mode);
+        cudaStreamSynchronize(device_cuda.stream);
+
+        cudaMemcpy(pred_buf.data(), gpu_output._data, BATCH_SIZE * OUTPUT_DIM * sizeof(float), cudaMemcpyDeviceToHost);
+
+        const float tan_half_hfov = std::tan(rlt::rendering::raytracing::degrees_to_radians(VIEWER_FOV) / 2.0f);
+
+        // ---- Composite 8x8 grid ----
+        std::memset(output_image.data(), 0, output_image.size());
+        for (TI i = 0; i < NUM_GRID_CAMERAS; i++) {
+            TI grid_col = i % GRID_DIM;
+            TI grid_row = i / GRID_DIM;
+            TI tile_x = grid_col * CAM_WIDTH;
+            TI tile_y = grid_row * CAM_HEIGHT;
+
+            const uint32_t* cam_pixels_ptr = tile_pixels[i].data();
+            for (TI y = 0; y < CAM_HEIGHT; y++) {
+                for (TI x = 0; x < CAM_WIDTH; x++) {
+                    uint32_t rgba = cam_pixels_ptr[y * CAM_WIDTH + x];
+                    TI px = tile_x + x;
+                    TI py = tile_y + y;
+                    TI idx = (py * IMAGE_WIDTH + px) * 3;
+                    output_image[idx + 0] = (rgba >>  0) & 0xFF;
+                    output_image[idx + 1] = (rgba >>  8) & 0xFF;
+                    output_image[idx + 2] = (rgba >> 16) & 0xFF;
+                }
+            }
+
+            float pred_angle = std::atan(pred_buf[i * OUTPUT_DIM + 0] * tan_half_hfov);
+
+            draw_bead(output_image, IMAGE_WIDTH, tile_x, tile_y, delta_yaw, MAX_ANGLE, 2, 0, 200, 0);
+            draw_bead(output_image, IMAGE_WIDTH, tile_x, tile_y, pred_angle, MAX_ANGLE, 5, 200, 0, 0);
+        }
+
+        // Write frame to ffmpeg pipe
+        size_t frame_bytes = output_image.size();
+        size_t written = fwrite(output_image.data(), 1, frame_bytes, mp4_pipe);
+        if (written != frame_bytes) {
+            std::cerr << "Failed writing frame " << frame << " to ffmpeg" << std::endl;
+            pclose(mp4_pipe);
+            mp4_pipe = nullptr;
+            break;
+        }
+
+        if (frame % 10 == 0 || frame == num_frames - 1) {
+            std::cout << "Frame " << frame << "/" << num_frames
+                      << "  delta_yaw=" << delta_yaw * 180.0f / PI << " deg" << std::endl;
+        }
+    }
+
+    if (mp4_pipe) {
+        int ffmpeg_status = pclose(mp4_pipe);
+        if (ffmpeg_status != 0) {
+            std::cerr << "ffmpeg exited with status " << ffmpeg_status << std::endl;
+        } else {
+            std::cout << "Video written: " << output_path << std::endl;
+        }
+    }
+
+    // ---- Cleanup ----
+    rlt::free(device_cuda, gpu_output);
+    rlt::free(device_cuda, gpu_input_a);
+    rlt::free(device_cuda, gpu_input_b);
+    rlt::free(device_cuda, model_buffer);
+    rlt::free(device_cuda, model);
+    rlt::free(device_cuda, rng_cuda);
+    rlt::free(device_cpu, model_cpu);
+    for (auto& s : loaded_scenes) {
+        yp::destroy_scene(s.handle);
+    }
+
+    return 0;
+}

hyperdrone/_vendor/rl-tools/src/rl/environments/hyperdrone/spy/CMakeLists.txt

@@ -0,0 +1,63 @@
+set(RL_TOOLS_ENABLE_HYPERDRONE_SPY OFF PARENT_SCOPE)
+if(RL_TOOLS_DISABLE_HYPERDRONE_SPY)
+    return()
+endif()
+if(NOT RL_TOOLS_RENDERING_ENABLE_RAYTRACING)
+    set(RL_TOOLS_SUMMARY_HYPERDRONE_SPY_REASON "raytracing unavailable" PARENT_SCOPE)
+    return()
+endif()
+
+find_package(glfw3 QUIET)
+find_package(OpenGL QUIET)
+if(NOT TARGET glfw OR NOT TARGET OpenGL::GL)
+    set(RL_TOOLS_SUMMARY_HYPERDRONE_SPY_REASON "GLFW/OpenGL not found (e.g. apt install libglfw3-dev libgl1-mesa-dev)" PARENT_SCOPE)
+    return()
+endif()
+
+set(RL_TOOLS_IMGUI_SOURCE_DIR "$ENV{RL_TOOLS_IMGUI_SOURCE_DIR}" CACHE PATH "Path to existing Dear ImGui sources for the HyperDrone spy")
+if(RL_TOOLS_IMGUI_SOURCE_DIR)
+    foreach(_source imgui.cpp imgui_draw.cpp imgui_tables.cpp imgui_widgets.cpp backends/imgui_impl_glfw.cpp backends/imgui_impl_opengl3.cpp)
+        if(NOT EXISTS "${RL_TOOLS_IMGUI_SOURCE_DIR}/${_source}")
+            message(FATAL_ERROR "RL_TOOLS_IMGUI_SOURCE_DIR is missing ${_source}")
+        endif()
+    endforeach()
+    set(_imgui_sources "${RL_TOOLS_IMGUI_SOURCE_DIR}/imgui.cpp" "${RL_TOOLS_IMGUI_SOURCE_DIR}/imgui_draw.cpp" "${RL_TOOLS_IMGUI_SOURCE_DIR}/imgui_tables.cpp" "${RL_TOOLS_IMGUI_SOURCE_DIR}/imgui_widgets.cpp")
+    set(_imgui_include_dirs "${RL_TOOLS_IMGUI_SOURCE_DIR}" "${RL_TOOLS_IMGUI_SOURCE_DIR}/backends")
+    set(_imgui_backend_dir "${RL_TOOLS_IMGUI_SOURCE_DIR}/backends")
+    set(_imgui_reason "Dear ImGui sources ${RL_TOOLS_IMGUI_SOURCE_DIR}")
+else()
+    find_package(PkgConfig QUIET)
+    if(PkgConfig_FOUND)
+        pkg_check_modules(RL_TOOLS_SPY_IMGUI QUIET IMPORTED_TARGET imgui>=1.86)
+    endif()
+    find_path(RL_TOOLS_IMGUI_BACKENDS_DIR NAMES imgui_impl_glfw.cpp HINTS ${RL_TOOLS_SPY_IMGUI_INCLUDE_DIRS} PATHS /usr/share/doc/libimgui-dev/examples /usr/local/share/imgui /usr/share/imgui PATH_SUFFIXES backends)
+    if(NOT TARGET PkgConfig::RL_TOOLS_SPY_IMGUI OR NOT EXISTS "${RL_TOOLS_IMGUI_BACKENDS_DIR}/imgui_impl_opengl3.cpp")
+        set(RL_TOOLS_SUMMARY_HYPERDRONE_SPY_REASON "Dear ImGui not found (e.g. apt install libimgui-dev, or set RL_TOOLS_IMGUI_SOURCE_DIR)" PARENT_SCOPE)
+        return()
+    endif()
+    find_path(RL_TOOLS_IMGUI_BACKENDS_INCLUDE_DIR NAMES imgui_impl_glfw.h HINTS ${RL_TOOLS_SPY_IMGUI_INCLUDE_DIRS} "${RL_TOOLS_IMGUI_BACKENDS_DIR}" PATH_SUFFIXES backends)
+    if(NOT RL_TOOLS_IMGUI_BACKENDS_INCLUDE_DIR)
+        set(RL_TOOLS_SUMMARY_HYPERDRONE_SPY_REASON "Dear ImGui backend headers not found" PARENT_SCOPE)
+        return()
+    endif()
+    set(_imgui_backend_dir "${RL_TOOLS_IMGUI_BACKENDS_DIR}")
+    set(_imgui_include_dirs "${RL_TOOLS_IMGUI_BACKENDS_INCLUDE_DIR}")
+    set(_imgui_reason "system Dear ImGui ${RL_TOOLS_SPY_IMGUI_VERSION}")
+endif()
+
+find_package(Threads REQUIRED)
+add_library(rl_tools_hyperdrone_spy SHARED ui.cpp ${_imgui_sources} "${_imgui_backend_dir}/imgui_impl_glfw.cpp" "${_imgui_backend_dir}/imgui_impl_opengl3.cpp")
+add_library(RLtools::HyperDroneSpy ALIAS rl_tools_hyperdrone_spy)
+target_link_libraries(rl_tools_hyperdrone_spy PUBLIC RLtools::Core Threads::Threads PRIVATE glfw OpenGL::GL ${CMAKE_DL_LIBS})
+if(TARGET PkgConfig::RL_TOOLS_SPY_IMGUI AND NOT RL_TOOLS_IMGUI_SOURCE_DIR)
+    target_link_libraries(rl_tools_hyperdrone_spy PRIVATE PkgConfig::RL_TOOLS_SPY_IMGUI)
+endif()
+target_include_directories(rl_tools_hyperdrone_spy SYSTEM PRIVATE ${_imgui_include_dirs})
+target_compile_definitions(rl_tools_hyperdrone_spy PUBLIC RL_TOOLS_ENABLE_HYPERDRONE_SPY=1)
+set_target_properties(rl_tools_hyperdrone_spy PROPERTIES CXX_STANDARD 17 CXX_STANDARD_REQUIRED ON WINDOWS_EXPORT_ALL_SYMBOLS ON)
+if(NOT RL_TOOLS_DISABLE_TARGETS)
+    add_executable(example_hyperdrone_spy example.cpp)
+    target_link_libraries(example_hyperdrone_spy PRIVATE RLtools::RLtools RLtools::HyperDroneSpy rendering_raytracing_backend rendering_raytracing_stb)
+endif()
+set(RL_TOOLS_ENABLE_HYPERDRONE_SPY ON PARENT_SCOPE)
+set(RL_TOOLS_SUMMARY_HYPERDRONE_SPY_REASON "GLFW/OpenGL, ${_imgui_reason}" PARENT_SCOPE)

hyperdrone/_vendor/rl-tools/src/rl/environments/hyperdrone/spy/example.cpp

@@ -0,0 +1,101 @@
+#include <rl_tools/operations/cpu_mux.h>
+#include <rl_tools/rl/environments/hyperdrone/presets.h>
+#include <rl_tools/rl/environments/hyperdrone/operations_cpu.h>
+#include <rl_tools/rl/environments/hyperdrone/spy/operations_cpu.h>
+#include <metra/metra.h>
+
+#include <cstdio>
+#include <exception>
+#include <filesystem>
+
+namespace rlt = rl_tools;
+namespace hd = rlt::rl::environments::hyperdrone;
+
+using DEVICE = rlt::devices::DEVICE_FACTORY<>;
+using T = float;
+using TI = DEVICE::index_t;
+using RNG = DEVICE::SPEC::RANDOM::ENGINE<>;
+constexpr bool DYNAMIC_ALLOCATION = true;
+constexpr TI SEED = 0;
+struct WORLD_SPEC: hd::presets::X500FPV<T, TI> {
+    static constexpr TI INSTANCES_PER_ENVIRONMENT = 4;
+    static constexpr TI CAM_WIDTH = 64;
+    static constexpr TI CAM_HEIGHT = 64;
+    static constexpr bool SELF_VISIBLE = false;
+};
+using WORLD = hd::World<WORLD_SPEC>;
+using ENVIRONMENT = hd::MultiEnvironment<WORLD, 1>;
+using SPY = hd::Spy<hd::spy::Config<ENVIRONMENT>>;
+template <typename ELEMENT, TI... DIMS>
+using InstanceTensor = rlt::Tensor<rlt::tensor::Specification<ELEMENT, TI, rlt::tensor::Shape<TI, ENVIRONMENT::INSTANCES, DIMS...>, DYNAMIC_ALLOCATION>>;
+
+int main(int argc, char** argv){
+    if(argc != 2 || !std::filesystem::is_regular_file(argv[1])){
+        std::fprintf(stderr, "usage: %s /path/to/local/scene.glb\n", argv[0]);
+        return 1;
+    }
+    using namespace rl_tools;
+    DEVICE device{};
+    ENVIRONMENT environment{};
+    SPY spy{};
+    RNG rng{};
+    InstanceTensor<WORLD::Parameters> parameters{};
+    InstanceTensor<WORLD::State> states{}, next_states{};
+    InstanceTensor<bool> reset_mask{};
+    InstanceTensor<T, WORLD::ACTION_DIM> actions{};
+    InstanceTensor<T, WORLD::OBSERVATION_DIM> observations{};
+
+    init(device);
+    malloc(device, environment);
+    malloc(device, rng);
+    malloc(device, parameters);
+    malloc(device, states);
+    malloc(device, next_states);
+    malloc(device, reset_mask);
+    malloc(device, actions);
+    malloc(device, observations);
+
+    int result = 0;
+    TI step_i = 0;
+    try{
+        rlt::rendering::datasets::glb::Dataset dataset{};
+        dataset.references = {argv[1]};
+        init(device, environment, dataset);
+        init(device, rng, SEED);
+        set_all(device, reset_mask, true);
+        set_all(device, actions, (T)0);
+        malloc(device, spy);
+        init(device, spy, environment);
+        while(true){
+            sample_initial_parameters(device, environment, parameters, reset_mask, rng);
+            sample_initial_state(device, environment, parameters, states, reset_mask, rng);
+            observe(device, environment, parameters, states, WORLD::Observation{}, observations, rng);
+            if(!update(device, spy, environment)){
+                break;
+            }
+            step(device, environment, parameters, states, actions, next_states, rng);
+            copy(device, device, next_states, states);
+            terminated(device, environment, parameters, states, reset_mask, rng);
+            if(++step_i % WORLD::EPISODE_STEP_LIMIT == 0){
+                set_all(device, reset_mask, true);
+            }
+        }
+    }
+    catch(const std::exception& error){
+        std::fprintf(stderr, "HyperDrone spy: %s\n", error.what());
+        result = 1;
+    }
+
+    free(device, spy);
+    free(device, observations);
+    free(device, actions);
+    free(device, reset_mask);
+    free(device, next_states);
+    free(device, states);
+    free(device, parameters);
+    free(device, rng);
+    free(device, environment);
+    free(device);
+    metra::log("hyperdrone/spy/example/steps", double(step_i));
+    return result;
+}