hyperray 0.0.1__py3-none-any.whl

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/environments/hyperdrone/rig/operations_cpu.h

@@ -0,0 +1,190 @@
+#include "../../../../version.h"
+#include "../../../../rl_tools.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ENVIRONMENTS_HYPERDRONE_RIG_OPERATIONS_CPU_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ENVIRONMENTS_HYPERDRONE_RIG_OPERATIONS_CPU_H
+
+#include "rig.h"
+#include "operations_generic.h"
+#include "../../../../rendering/rig/operations_cpu.h"
+#include "../../../../rendering/raytracing/operations_cpu_common.h"
+
+#include <cstring>
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools {
+    namespace rl::environments::hyperdrone::rig::detail {
+        inline bool part_bounds(const rendering::Object& object, float low[3], float high[3]){
+            bool first = true;
+            for(const auto& mesh : object.meshes){
+                for(size_t vertex_i = 0; vertex_i + 2 < mesh.vertices.size(); vertex_i += 3){
+                    for(unsigned dim = 0; dim < 3; dim++){
+                        const float value = mesh.vertices[vertex_i + dim];
+                        low[dim] = first || value < low[dim] ? value : low[dim];
+                        high[dim] = first || value > high[dim] ? value : high[dim];
+                    }
+                    first = false;
+                }
+            }
+            return !first;
+        }
+
+        template <typename TI, TI MAX_PROPS>
+        struct Props {
+            TI count = 0;
+            TI parts[MAX_PROPS];
+            float pivots[MAX_PROPS][3]; // hub in the part frame
+            float hubs[MAX_PROPS][3];   // hub in the assembly frame
+        };
+
+        // validates the body/prop_* convention (the first part is "body", a prop_* part has geometry)
+        template <typename TI, TI MAX_PROPS>
+        bool collect_props(const rendering::ObjectAssembly& assembly, TI max_parts, Props<TI, MAX_PROPS>& props){
+            if(assembly.parts.empty() || assembly.parts.size() > max_parts || assembly.objects[assembly.parts[0].object].name != "body"){
+                return false;
+            }
+            props.count = 0;
+            for(TI part_i = 0; part_i < (TI)assembly.parts.size(); part_i++){
+                const auto& part = assembly.parts[part_i];
+                const auto& object = assembly.objects[part.object];
+                if(object.name.rfind("prop_", 0) == 0){
+                    float low[3] = {0, 0, 0}, high[3] = {0, 0, 0};
+                    if(props.count >= MAX_PROPS || !part_bounds(object, low, high)){
+                        return false;
+                    }
+                    float* pivot = props.pivots[props.count];
+                    pivot[0] = (low[0] + high[0]) / 2;
+                    pivot[1] = (low[1] + high[1]) / 2;
+                    pivot[2] = 0;
+                    rendering::transform_point(part.transform, pivot, props.hubs[props.count]);
+                    props.parts[props.count] = part_i;
+                    props.count++;
+                }
+                else if(object.name == "body" && part_i != 0){
+                    return false;
+                }
+            }
+            return true;
+        }
+
+        // prop p becomes link 1 + coordinate_of_prop[p], revolute about its hub along the part's z
+        // with the sign spin[p]; every other part rides link 0 at its rest frame
+        template <typename DEVICE, typename TI, TI MAX_PROPS, TI MAX_PARTS>
+        bool assemble(DEVICE& device, Rotorcraft<TI, MAX_PROPS, MAX_PARTS>& rig, const rendering::ObjectAssembly& assembly, const Props<TI, MAX_PROPS>& props, const TI* coordinate_of_prop, const float* spin){
+            using SPEC = typename Rotorcraft<TI, MAX_PROPS, MAX_PARTS>::SPEC;
+            rig.num_links = 1 + props.count;
+            rig.num_parts = (TI)assembly.parts.size();
+            rig.links[0] = rendering::rig::Link<SPEC>{};
+            for(TI part_i = 0; part_i < rig.num_parts; part_i++){
+                rig.parts[part_i] = rendering::rig::Part<SPEC>{};
+                rig.parts[part_i].link = 0;
+                std::memcpy(rig.parts[part_i].offset, assembly.parts[part_i].transform, sizeof(rig.parts[part_i].offset));
+            }
+            for(TI prop_i = 0; prop_i < props.count; prop_i++){
+                const TI link_i = 1 + coordinate_of_prop[prop_i];
+                const float* pivot = props.pivots[prop_i];
+                const float to_pivot[12] = {1, 0, 0, pivot[0], 0, 1, 0, pivot[1], 0, 0, 1, pivot[2]};
+                auto& link = rig.links[link_i];
+                link = rendering::rig::Link<SPEC>{};
+                link.parent = 0;
+                link.joint = rendering::rig::Joint::REVOLUTE;
+                link.axis[0] = 0;
+                link.axis[1] = 0;
+                link.axis[2] = spin[prop_i];
+                rendering::compose_transforms(assembly.parts[props.parts[prop_i]].transform, to_pivot, link.origin);
+                auto& part = rig.parts[props.parts[prop_i]];
+                part = rendering::rig::Part<SPEC>{};
+                part.link = link_i;
+                for(unsigned axis = 0; axis < 3; axis++){
+                    part.offset[axis * 4 + 3] = -pivot[axis];
+                }
+            }
+            return init(device, static_cast<rendering::rig::Model<SPEC>&>(rig));
+        }
+    }
+
+    // "body"/"prop_*" part-name convention -> rig: coordinate i is the phase of the i-th prop in
+    // part order, spinning per the Quad-X convention (front-right CCW) by its hub's quadrant in
+    // the assembly frame; false on convention violations (fail-fast for rigged-drone datasets)
+    template <typename DEVICE, typename TI, TI MAX_PROPS, TI MAX_PARTS>
+    bool init(DEVICE& device, rl::environments::hyperdrone::rig::Rotorcraft<TI, MAX_PROPS, MAX_PARTS>& rig, const rendering::ObjectAssembly& assembly){
+        rl::environments::hyperdrone::rig::detail::Props<TI, MAX_PROPS> props;
+        if(!rl::environments::hyperdrone::rig::detail::collect_props(assembly, MAX_PARTS, props)){
+            return false;
+        }
+        TI coordinate_of_prop[MAX_PROPS];
+        float spin[MAX_PROPS];
+        for(TI prop_i = 0; prop_i < props.count; prop_i++){
+            coordinate_of_prop[prop_i] = prop_i;
+            spin[prop_i] = props.hubs[prop_i][0] * props.hubs[prop_i][1] < 0 ? 1.0f : -1.0f;
+        }
+        return rl::environments::hyperdrone::rig::detail::assemble(device, rig, assembly, props, coordinate_of_prop, spin);
+    }
+
+    // the same convention matched to a dynamics model: each hub pairs with the nearest rotor in
+    // the assembly XY plane (a bijection, else false), coordinate i is the phase of rotor i and
+    // the prop spins against the rotor's reaction torque
+    template <typename DEVICE, typename TI, TI MAX_PROPS, TI MAX_PARTS, typename T>
+    bool init(DEVICE& device, rl::environments::hyperdrone::rig::Rotorcraft<TI, MAX_PROPS, MAX_PARTS>& rig, const rendering::ObjectAssembly& assembly, const T rotor_positions[][3], const T rotor_torque_directions[][3], TI num_rotors){
+        rl::environments::hyperdrone::rig::detail::Props<TI, MAX_PROPS> props;
+        if(!rl::environments::hyperdrone::rig::detail::collect_props(assembly, MAX_PARTS, props) || props.count != num_rotors){
+            return false;
+        }
+        TI coordinate_of_prop[MAX_PROPS];
+        float spin[MAX_PROPS];
+        bool matched[MAX_PROPS] = {};
+        for(TI prop_i = 0; prop_i < props.count; prop_i++){
+            TI nearest = 0;
+            T nearest_distance = 0;
+            for(TI rotor_i = 0; rotor_i < num_rotors; rotor_i++){
+                const T dx = (T)props.hubs[prop_i][0] - rotor_positions[rotor_i][0];
+                const T dy = (T)props.hubs[prop_i][1] - rotor_positions[rotor_i][1];
+                const T distance = dx * dx + dy * dy;
+                if(rotor_i == 0 || distance < nearest_distance){
+                    nearest = rotor_i;
+                    nearest_distance = distance;
+                }
+            }
+            const T torque = rotor_torque_directions[nearest][2];
+            if(matched[nearest] || torque == 0){
+                return false;
+            }
+            matched[nearest] = true;
+            coordinate_of_prop[prop_i] = nearest;
+            spin[prop_i] = torque > 0 ? -1.0f : 1.0f;
+        }
+        return rl::environments::hyperdrone::rig::detail::assemble(device, rig, assembly, props, coordinate_of_prop, spin);
+    }
+
+    // the rig verbs: evaluate/sample the rig at the base pose and rig.num_coordinates joint
+    // values (for a Rotorcraft the prop phases in coordinate order), then publish the part poses
+    template <typename DEVICE, typename SPEC, typename BACKEND, typename RIG_SPEC, typename T>
+    void set_transform(DEVICE& device, rendering::raytracing::Renderer<SPEC, BACKEND>& renderer, rendering::raytracing::OverlayIndex overlay, const rendering::raytracing::OverlayPlacement& placement, const rendering::rig::Model<RIG_SPEC>& rig, const float base[12], const T* coordinates){
+        using TI = typename RIG_SPEC::TI;
+        float coordinates_float[RIG_SPEC::MAX_COORDINATES];
+        for(TI coordinate_i = 0; coordinate_i < rig.num_coordinates; coordinate_i++){
+            coordinates_float[coordinate_i] = (float)coordinates[coordinate_i];
+        }
+        rendering::rig::Buffer<RIG_SPEC> buffer;
+        float poses[RIG_SPEC::MAX_PARTS][12];
+        evaluate(device, rig, base, coordinates_float, buffer, &poses[0][0]);
+        set_transforms(device, renderer, overlay, placement, &poses[0][0]);
+    }
+
+    template <typename DEVICE, typename SPEC, typename BACKEND, typename RIG_SPEC, typename T>
+    void set_transform_pair(DEVICE& device, rendering::raytracing::Renderer<SPEC, BACKEND>& renderer, rendering::raytracing::OverlayIndex overlay, const rendering::raytracing::OverlayPlacement& placement, const rendering::rig::Model<RIG_SPEC>& rig, const float base_open[12], const float base_close[12], const T* coordinates_open, const T* coordinates_close){
+        using TI = typename RIG_SPEC::TI;
+        float coordinates_open_float[RIG_SPEC::MAX_COORDINATES], coordinates_close_float[RIG_SPEC::MAX_COORDINATES];
+        for(TI coordinate_i = 0; coordinate_i < rig.num_coordinates; coordinate_i++){
+            coordinates_open_float[coordinate_i] = (float)coordinates_open[coordinate_i];
+            coordinates_close_float[coordinate_i] = (float)coordinates_close[coordinate_i];
+        }
+        rendering::rig::Buffer<RIG_SPEC> buffer;
+        float samples[SPEC::MOTION_BLUR_SAMPLES][RIG_SPEC::MAX_PARTS][12], close[RIG_SPEC::MAX_PARTS][12];
+        sample(device, rig, base_open, base_close, coordinates_open_float, coordinates_close_float, (TI)SPEC::MOTION_BLUR_SAMPLES, buffer, &samples[0][0][0], (TI)(rig.num_parts * 12), &close[0][0]);
+        set_motion_transforms(device, renderer, overlay, placement, &samples[0][0][0], &close[0][0]);
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/environments/hyperdrone/rig/operations_generic.h

@@ -0,0 +1,32 @@
+#include "../../../../version.h"
+#include "../../../../rl_tools.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ENVIRONMENTS_HYPERDRONE_RIG_OPERATIONS_GENERIC_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ENVIRONMENTS_HYPERDRONE_RIG_OPERATIONS_GENERIC_H
+
+#include "rig.h"
+#include "../../../../rendering/rig/operations_generic.h"
+#include "../../l2f/quaternion_helper.h"
+
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::rl::environments::hyperdrone::rig {
+    template <typename DEVICE, typename T>
+    RL_TOOLS_FUNCTION_PLACEMENT void make_body_transform(DEVICE&, const T orientation[4], const T position[3], const T scene_translation[3], T scene_yaw_cos, T scene_yaw_sin, float out[12]){
+        T R[3][3];
+        rl::environments::l2f::quaternion_to_rotation_matrix<DEVICE, T>(orientation, R);
+        const T yc = scene_yaw_cos;
+        const T ys = scene_yaw_sin;
+        for(unsigned column = 0; column < 3; column++){
+            out[0 * 4 + column] = (float)(yc * R[0][column] - ys * R[1][column]);
+            out[1 * 4 + column] = (float)(ys * R[0][column] + yc * R[1][column]);
+            out[2 * 4 + column] = (float)R[2][column];
+        }
+        out[0 * 4 + 3] = (float)(yc * position[0] - ys * position[1] + scene_translation[0]);
+        out[1 * 4 + 3] = (float)(ys * position[0] + yc * position[1] + scene_translation[1]);
+        out[2 * 4 + 3] = (float)(position[2] + scene_translation[2]);
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/environments/hyperdrone/rig/rig.h

@@ -0,0 +1,25 @@
+#include "../../../../version.h"
+#include "../../../../rl_tools.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ENVIRONMENTS_HYPERDRONE_RIG_RIG_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ENVIRONMENTS_HYPERDRONE_RIG_RIG_H
+
+#include "../../../../rendering/rig/rig.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::rl::environments::hyperdrone::rig {
+    // the body/prop_* assembly convention as a rendering rig: link 0 is the assembly frame (fixed
+    // at identity), each prop_* part a revolute link about its hub (the part's AABB centre, since
+    // node origins need not sit at the rotor hubs in real assets), every other part rides link 0 at
+    // its rest frame. Built from the convention alone, coordinate i is the phase of the i-th prop
+    // in part order and the spin follows the Quad-X convention (front-right CCW) from the hub's
+    // quadrant; built against a dynamics model, coordinate i is the phase of rotor i (hubs matched
+    // to the nearest rotor) spinning against the rotor's reaction torque.
+    template <typename T_TI, T_TI T_MAX_PROPS = 4, T_TI T_MAX_PARTS = 8>
+    struct Rotorcraft: rendering::rig::Model<rendering::rig::Specification<float, T_TI, 1 + T_MAX_PROPS, T_MAX_PARTS>> {
+        static constexpr T_TI MAX_PROPS = T_MAX_PROPS;
+    };
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/environments/hyperdrone/spy/capture_cpu.h

@@ -0,0 +1,86 @@
+#include "spy.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ENVIRONMENTS_HYPERDRONE_SPY_CAPTURE_CPU_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ENVIRONMENTS_HYPERDRONE_SPY_CAPTURE_CPU_H
+
+#include <algorithm>
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools {
+    template <typename CONFIG>
+    void pack_spy_atlas_cpu(rl::environments::hyperdrone::spy::State<CONFIG>& state, typename CONFIG::ENVIRONMENT& environment) {
+        using WORLD = typename CONFIG::WORLD;
+        constexpr size_t CAMERAS = WORLD::INSTANCES * WORLD::N_VIEWS;
+        constexpr size_t WIDTH = CONFIG::ATLAS_WIDTH, HEIGHT = CONFIG::ATLAS_HEIGHT;
+        for(size_t member = 0; member < CONFIG::ENVIRONMENT::NUMBER_OF_ENVIRONMENTS; member++) {
+            auto& world = environment.environments[member];
+            const bool valid = state.buffer.frames[member].valid;
+            const float* history = valid ? data(world.history) + ((world.history_step - 1) % get<0>(typename WORLD::HISTORY_SPEC::SHAPE{})) * CAMERAS * WORLD::FRAME_DIM : nullptr;
+            for(size_t camera = 0; camera < CAMERAS; camera++) {
+                const size_t index = member * CAMERAS + camera;
+                const size_t column = index % state.columns, row = index / state.columns;
+                const size_t x0 = (column * WIDTH + state.columns - 1) / state.columns, x1 = ((column + 1) * WIDTH + state.columns - 1) / state.columns;
+                const size_t y0 = (row * HEIGHT + state.rows - 1) / state.rows, y1 = ((row + 1) * HEIGHT + state.rows - 1) / state.rows;
+                for(size_t y = y0; y < y1; y++) {
+                    for(size_t x = x0; x < x1; x++) {
+                        uint32_t pixel = 0xff202020;
+                        if(valid) {
+                            const size_t source_x = (x - x0) * WORLD::RENDERER_CONFIG::CAM_WIDTH / (x1 - x0);
+                            const size_t source_y = (y - y0) * WORLD::RENDERER_CONFIG::CAM_HEIGHT / (y1 - y0);
+                            const float* rgb = history + camera * WORLD::FRAME_DIM + (source_y * WORLD::RENDERER_CONFIG::CAM_WIDTH + source_x) * WORLD::IMAGE_CHANNELS;
+                            pixel = 0xff000000;
+                            for(size_t channel = 0; channel < 3; channel++) {
+                                const float value = std::max(0.0f, std::min(1.0f, rgb[channel]));
+                                pixel |= static_cast<uint32_t>(value * 255.0f + 0.5f) << (channel * 8);
+                            }
+                        }
+                        state.buffer.atlas[y * WIDTH + x] = pixel;
+                    }
+                }
+            }
+        }
+    }
+
+    template <typename DEVICE, typename CONFIG, typename BACKEND>
+    void malloc_spy_capture(DEVICE&, rl::environments::hyperdrone::spy::State<CONFIG>& state, BACKEND) {
+        state.buffer.atlas = new uint32_t[CONFIG::ATLAS_WIDTH * CONFIG::ATLAS_HEIGHT];
+        std::fill_n(state.buffer.atlas, CONFIG::ATLAS_WIDTH * CONFIG::ATLAS_HEIGHT, 0xff202020);
+    }
+
+    template <typename DEVICE, typename CONFIG, typename BACKEND>
+    void free_spy_capture(DEVICE&, rl::environments::hyperdrone::spy::State<CONFIG>& state, BACKEND) {
+        delete[] state.buffer.atlas;
+        state.buffer.atlas = nullptr;
+    }
+
+    template <typename DEVICE, typename CONFIG, typename BACKEND>
+    bool poll_spy_capture(DEVICE&, rl::environments::hyperdrone::spy::State<CONFIG>&, BACKEND) {
+        return true;
+    }
+
+    template <typename DEVICE, typename CONFIG, typename BACKEND>
+    void capture_spy(DEVICE& device, rl::environments::hyperdrone::spy::State<CONFIG>& state, typename CONFIG::ENVIRONMENT& environment, BACKEND) {
+        using WORLD = typename CONFIG::WORLD;
+        pack_spy_atlas_cpu(state, environment);
+        if(state.buffer.selected_valid) {
+            const size_t camera = state.buffer.selected_camera % (WORLD::INSTANCES * WORLD::N_VIEWS);
+            const size_t member = state.buffer.selected_camera / (WORLD::INSTANCES * WORLD::N_VIEWS);
+            const size_t instance = camera / WORLD::N_VIEWS;
+            auto& world = environment.environments[member];
+            auto& renderer = world.slots[world.active_slot].renderer;
+            state.buffer.geometry.camera = data(cameras(device, renderer))[camera];
+            if constexpr(WORLD::RENDERER_SPEC::ENABLE_OVERLAYS) {
+                for(size_t slot = 0; slot < WORLD::RENDERER_SPEC::MAX_OVERLAY_INSTANCES; slot++) {
+                    auto& destination = state.buffer.geometry.slots[slot];
+                    const auto& source = renderer.overlays[instance].slots[slot];
+                    destination.object = source.object;
+                    destination.active = source.active;
+                    const float* transform = data(transforms(device, renderer)) + (instance * WORLD::RENDERER_SPEC::MAX_OVERLAY_INSTANCES + slot) * 12;
+                    std::copy_n(transform, 12, destination.transform);
+                }
+            }
+        }
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/environments/hyperdrone/spy/capture_cuda.h

@@ -0,0 +1,69 @@
+#include "../../../../version.h"
+#include "../../../../rl_tools.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ENVIRONMENTS_HYPERDRONE_SPY_CAPTURE_CUDA_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ENVIRONMENTS_HYPERDRONE_SPY_CAPTURE_CUDA_H
+
+#include "capture_optix.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::rl::environments::hyperdrone::spy {
+    template <typename CONFIG>
+    __global__ void pack_spy_atlas_kernel(const float* history, uint32_t* atlas, size_t member, size_t columns, size_t rows, size_t first_y, size_t pixels) {
+        using WORLD = typename CONFIG::WORLD;
+        constexpr size_t CAMERAS = WORLD::INSTANCES * WORLD::N_VIEWS;
+        constexpr size_t WIDTH = CONFIG::ATLAS_WIDTH, HEIGHT = CONFIG::ATLAS_HEIGHT;
+        const size_t pixel_i = static_cast<size_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+        if(pixel_i < pixels) {
+            const size_t x = pixel_i % WIDTH, y = first_y + pixel_i / WIDTH;
+            const size_t column = x * columns / WIDTH, row = y * rows / HEIGHT;
+            const size_t camera = row * columns + column;
+            if(camera >= member * CAMERAS && camera < (member + 1) * CAMERAS) {
+                uint32_t pixel = 0xff202020;
+                if(history != nullptr) {
+                    const size_t x0 = (column * WIDTH + columns - 1) / columns, x1 = ((column + 1) * WIDTH + columns - 1) / columns;
+                    const size_t y0 = (row * HEIGHT + rows - 1) / rows, y1 = ((row + 1) * HEIGHT + rows - 1) / rows;
+                    const size_t source_x = (x - x0) * WORLD::RENDERER_CONFIG::CAM_WIDTH / (x1 - x0);
+                    const size_t source_y = (y - y0) * WORLD::RENDERER_CONFIG::CAM_HEIGHT / (y1 - y0);
+                    const float* rgb = history + (camera - member * CAMERAS) * WORLD::FRAME_DIM + (source_y * WORLD::RENDERER_CONFIG::CAM_WIDTH + source_x) * WORLD::IMAGE_CHANNELS;
+                    pixel = 0xff000000;
+                    for(size_t channel = 0; channel < 3; channel++) {
+                        float value = rgb[channel] < 1.0f ? rgb[channel] : 1.0f;
+                        value = value > 0.0f ? value : 0.0f;
+                        pixel |= static_cast<uint32_t>(__fadd_rn(__fmul_rn(value, 255.0f), 0.5f)) << (channel * 8);
+                    }
+                }
+                atlas[y * WIDTH + x] = pixel;
+            }
+        }
+    }
+}
+
+namespace rl_tools {
+    template <typename DEV_SPEC, typename CONFIG>
+    void capture_spy(devices::CUDA<DEV_SPEC>& device, rl::environments::hyperdrone::spy::State<CONFIG>& state, typename CONFIG::ENVIRONMENT& env, rendering::raytracing::backends::Optix) {
+        using WORLD = typename CONFIG::WORLD;
+        constexpr size_t CAMERAS = WORLD::INSTANCES * WORLD::N_VIEWS;
+        constexpr size_t WIDTH = CONFIG::ATLAS_WIDTH, HEIGHT = CONFIG::ATLAS_HEIGHT;
+        auto& capture = *state.buffer.capture.state;
+        capture.producer_stream = device.stream;
+        capture.capturing = true;
+        for(size_t member = 0; member < CONFIG::ENVIRONMENT::NUMBER_OF_ENVIRONMENTS; member++) {
+            const auto& world = env.environments[member];
+            const float* history = state.buffer.frames[member].valid ? data(world.history) + ((world.history_step - 1) % get<0>(typename WORLD::HISTORY_SPEC::SHAPE{})) * CAMERAS * WORLD::FRAME_DIM : nullptr;
+            const size_t first_row = member * CAMERAS / state.columns, last_row = ((member + 1) * CAMERAS - 1) / state.columns;
+            const size_t y0 = (first_row * HEIGHT + state.rows - 1) / state.rows, y1 = ((last_row + 1) * HEIGHT + state.rows - 1) / state.rows;
+            const size_t pixels = WIDTH * (y1 - y0);
+            if(pixels > 0) {
+                constexpr size_t BLOCK_SIZE = 256;
+                rl::environments::hyperdrone::spy::pack_spy_atlas_kernel<CONFIG><<<(pixels + BLOCK_SIZE - 1) / BLOCK_SIZE, BLOCK_SIZE, 0, device.stream>>>(history, capture.device_atlas, member, state.columns, state.rows, y0, pixels);
+                spy_capture_cuda_check(cudaGetLastError());
+            }
+        }
+        capture_spy_geometry(device, state, env, device.stream);
+        transfer_spy_capture(device, state, true);
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/environments/hyperdrone/spy/capture_optix.h

@@ -0,0 +1,163 @@
+#include "../../../../version.h"
+#include "../../../../rl_tools.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ENVIRONMENTS_HYPERDRONE_SPY_CAPTURE_OPTIX_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ENVIRONMENTS_HYPERDRONE_SPY_CAPTURE_OPTIX_H
+
+#include "spy.h"
+#include "../../../../rendering/raytracing/backends/optix/operations_cuda.h"
+
+#include <algorithm>
+#include <cstring>
+#include <stdexcept>
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::rl::environments::hyperdrone::spy {
+    struct CaptureStorage<rendering::raytracing::backends::Optix>::State {
+        uint32_t* device_atlas = nullptr;
+        void* device_geometry = nullptr;
+        void* host_geometry = nullptr;
+        cudaStream_t copy_stream = nullptr;
+        cudaStream_t producer_stream = nullptr;
+        cudaEvent_t packed = nullptr;
+        cudaEvent_t ready = nullptr;
+        bool capturing = false;
+    };
+
+    template <typename CONFIG>
+    struct OptixGeometry {
+        static constexpr size_t SLOTS = Geometry<CONFIG>::SLOTS;
+        rendering::raytracing::Camera<typename CONFIG::T> camera;
+        rendering::raytracing::backends::optix::OverlaySlotStructure structure[SLOTS > 0 ? SLOTS : 1];
+        float transforms[SLOTS > 0 ? SLOTS : 1][12];
+    };
+}
+
+namespace rl_tools {
+    inline void spy_capture_cuda_check(cudaError_t status) {
+        if(status != cudaSuccess) {
+            throw std::runtime_error(std::string("HyperDrone spy capture: ") + cudaGetErrorString(status));
+        }
+    }
+
+    template <typename DEVICE, typename CONFIG>
+    void malloc_spy_capture(DEVICE&, rl::environments::hyperdrone::spy::State<CONFIG>& state, rendering::raytracing::backends::Optix) {
+        namespace spy = rl::environments::hyperdrone::spy;
+        auto* capture = new spy::CaptureStorage<rendering::raytracing::backends::Optix>::State{};
+        state.buffer.capture.state = capture;
+        const size_t atlas_bytes = CONFIG::ATLAS_WIDTH * CONFIG::ATLAS_HEIGHT * sizeof(uint32_t);
+        spy_capture_cuda_check(cudaMallocHost(&state.buffer.atlas, atlas_bytes));
+        spy_capture_cuda_check(cudaMalloc(&capture->device_atlas, atlas_bytes));
+        std::fill_n(state.buffer.atlas, CONFIG::ATLAS_WIDTH * CONFIG::ATLAS_HEIGHT, 0xff202020);
+        spy_capture_cuda_check(cudaMemcpy(capture->device_atlas, state.buffer.atlas, atlas_bytes, cudaMemcpyHostToDevice));
+        spy_capture_cuda_check(cudaMalloc(&capture->device_geometry, sizeof(spy::OptixGeometry<CONFIG>)));
+        spy_capture_cuda_check(cudaMallocHost(&capture->host_geometry, sizeof(spy::OptixGeometry<CONFIG>)));
+        spy_capture_cuda_check(cudaStreamCreateWithFlags(&capture->copy_stream, cudaStreamNonBlocking));
+        spy_capture_cuda_check(cudaEventCreateWithFlags(&capture->packed, cudaEventDisableTiming));
+        spy_capture_cuda_check(cudaEventCreateWithFlags(&capture->ready, cudaEventDisableTiming));
+    }
+
+    template <typename DEVICE, typename CONFIG>
+    void free_spy_capture(DEVICE&, rl::environments::hyperdrone::spy::State<CONFIG>& state, rendering::raytracing::backends::Optix) {
+        auto* capture = state.buffer.capture.state;
+        if(capture != nullptr) {
+            if(capture->capturing) {
+                cudaStreamSynchronize(capture->producer_stream);
+                cudaStreamSynchronize(capture->copy_stream);
+            }
+            if(capture->packed != nullptr) {
+                cudaEventDestroy(capture->packed);
+            }
+            if(capture->ready != nullptr) {
+                cudaEventDestroy(capture->ready);
+            }
+            if(capture->copy_stream != nullptr) {
+                cudaStreamDestroy(capture->copy_stream);
+            }
+            cudaFree(capture->device_atlas);
+            cudaFree(capture->device_geometry);
+            cudaFreeHost(capture->host_geometry);
+            delete capture;
+            state.buffer.capture.state = nullptr;
+        }
+        cudaFreeHost(state.buffer.atlas);
+        state.buffer.atlas = nullptr;
+    }
+
+    template <typename DEVICE, typename CONFIG>
+    void capture_spy_geometry(DEVICE& device, rl::environments::hyperdrone::spy::State<CONFIG>& state, typename CONFIG::ENVIRONMENT& env, cudaStream_t producer_stream) {
+        using WORLD = typename CONFIG::WORLD;
+        using GEOMETRY = rl::environments::hyperdrone::spy::OptixGeometry<CONFIG>;
+        auto& capture = *state.buffer.capture.state;
+        capture.producer_stream = producer_stream;
+        capture.capturing = true;
+        if(state.buffer.selected_valid) {
+            constexpr size_t CAMERAS_PER_MEMBER = WORLD::INSTANCES * WORLD::N_VIEWS;
+            const size_t member_i = state.buffer.selected_camera / CAMERAS_PER_MEMBER;
+            const size_t local_camera = state.buffer.selected_camera % CAMERAS_PER_MEMBER;
+            const size_t instance_i = local_camera / WORLD::N_VIEWS;
+            auto& world = env.environments[member_i];
+            auto& renderer = world.slots[world.active_slot].renderer;
+            auto* geometry = static_cast<GEOMETRY*>(capture.device_geometry);
+            spy_capture_cuda_check(cudaMemcpyAsync(&geometry->camera, data(cameras(device, renderer)) + local_camera, sizeof(geometry->camera), cudaMemcpyDeviceToDevice, producer_stream));
+            if constexpr(WORLD::RENDERER_SPEC::ENABLE_OVERLAYS) {
+                spy_capture_cuda_check(cudaMemcpyAsync(geometry->structure, overlay_structure(device, renderer) + instance_i * GEOMETRY::SLOTS, sizeof(geometry->structure), cudaMemcpyDeviceToDevice, producer_stream));
+                spy_capture_cuda_check(cudaMemcpyAsync(geometry->transforms, data(transforms(device, renderer)) + instance_i * GEOMETRY::SLOTS * 12, sizeof(geometry->transforms), cudaMemcpyDeviceToDevice, producer_stream));
+            }
+        }
+    }
+
+    template <typename DEVICE, typename CONFIG>
+    void transfer_spy_capture(DEVICE&, rl::environments::hyperdrone::spy::State<CONFIG>& state, bool transfer_atlas) {
+        auto& capture = *state.buffer.capture.state;
+        spy_capture_cuda_check(cudaEventRecord(capture.packed, capture.producer_stream));
+        spy_capture_cuda_check(cudaStreamWaitEvent(capture.copy_stream, capture.packed, 0));
+        if(transfer_atlas) {
+            const size_t atlas_bytes = CONFIG::ATLAS_WIDTH * CONFIG::ATLAS_HEIGHT * sizeof(uint32_t);
+            spy_capture_cuda_check(cudaMemcpyAsync(state.buffer.atlas, capture.device_atlas, atlas_bytes, cudaMemcpyDeviceToHost, capture.copy_stream));
+        }
+        if(state.buffer.selected_valid) {
+            spy_capture_cuda_check(cudaMemcpyAsync(capture.host_geometry, capture.device_geometry, sizeof(rl::environments::hyperdrone::spy::OptixGeometry<CONFIG>), cudaMemcpyDeviceToHost, capture.copy_stream));
+        }
+        spy_capture_cuda_check(cudaEventRecord(capture.ready, capture.copy_stream));
+        state.buffer.pending = true;
+    }
+
+    template <typename DEVICE, typename CONFIG>
+    bool poll_spy_capture(DEVICE&, rl::environments::hyperdrone::spy::State<CONFIG>& state, rendering::raytracing::backends::Optix) {
+        auto& capture = *state.buffer.capture.state;
+        const auto status = cudaEventQuery(capture.ready);
+        if(status == cudaErrorNotReady) {
+            return false;
+        }
+        spy_capture_cuda_check(status);
+        if(state.buffer.selected_valid) {
+            const auto& geometry = *static_cast<const rl::environments::hyperdrone::spy::OptixGeometry<CONFIG>*>(capture.host_geometry);
+            state.buffer.geometry.camera = geometry.camera;
+            if constexpr(CONFIG::WORLD::RENDERER_SPEC::ENABLE_OVERLAYS) {
+                for(size_t slot_i = 0; slot_i < geometry.SLOTS; slot_i++) {
+                    auto& destination = state.buffer.geometry.slots[slot_i];
+                    destination.active = geometry.structure[slot_i].active != 0;
+                    destination.object = geometry.structure[slot_i].object;
+                    std::memcpy(destination.transform, geometry.transforms[slot_i], sizeof(destination.transform));
+                }
+            }
+        }
+        capture.capturing = false;
+        return true;
+    }
+
+    template <typename DEV_SPEC, typename CONFIG>
+    void capture_spy(devices::CPU<DEV_SPEC>& device, rl::environments::hyperdrone::spy::State<CONFIG>& state, typename CONFIG::ENVIRONMENT& env, rendering::raytracing::backends::Optix) {
+        pack_spy_atlas_cpu(state, env);
+        constexpr size_t CAMERAS_PER_MEMBER = CONFIG::WORLD::INSTANCES * CONFIG::WORLD::N_VIEWS;
+        const size_t member_i = state.buffer.selected_camera / CAMERAS_PER_MEMBER;
+        auto& world = env.environments[member_i];
+        auto& renderer = world.slots[world.active_slot].renderer;
+        capture_spy_geometry(device, state, env, stream(device, renderer));
+        transfer_spy_capture(device, state, false);
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif