hyperverse 0.0.1__py3-none-any.whl

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

@@ -0,0 +1,224 @@
+#include "../../../../../version.h"
+#include "../../../../../rl_tools.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ENVIRONMENTS_HYPERDRONE_TASKS_TARGET_FRAME_OPERATIONS_CPU_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ENVIRONMENTS_HYPERDRONE_TASKS_TARGET_FRAME_OPERATIONS_CPU_H
+
+#include "target_frame.h"
+#include "../../operations_cpu.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools {
+
+    namespace rl::environments::hyperdrone::tasks::target_frame {
+        template <typename DEVICE, typename TASK_SPEC, typename DRONE_SPEC, typename RNG>
+        RL_TOOLS_FUNCTION_PLACEMENT void _sample_initial_parameters(DEVICE& device, typename World<TASK_SPEC>::NEXT_WORLD::DYNAMICS_ENV& dynamics, const typename World<TASK_SPEC>::NEXT_WORLD::Parameters& defaults, Tensor<DRONE_SPEC>& drones, typename TASK_SPEC::TI num_drones, typename World<TASK_SPEC>::Parameters& parameters, RNG& rng){
+            using T = typename TASK_SPEC::T;
+            using NEXT_WORLD = typename TASK_SPEC::NEXT_WORLD;
+            hyperdrone::_sample_initial_parameters<DEVICE, typename NEXT_WORLD::SPEC>(device, dynamics, defaults, drones, num_drones, static_cast<typename NEXT_WORLD::Parameters&>(parameters), rng);
+            parameters.target_pose[0] = 0;
+            parameters.target_pose[1] = 0;
+            parameters.target_pose[2] = 0;
+            parameters.target_pose[3] = 1;
+            parameters.target_pose[4] = 0;
+            parameters.target_pose[5] = 0;
+            parameters.target_pose[6] = 0;
+            if constexpr (TASK_SPEC::TARGET_FRAME_ROLL_PITCH_RANDOMIZATION_RANGE > (T)0){
+                parameters.target_roll = random::uniform_real_distribution(device.random, -TASK_SPEC::TARGET_FRAME_ROLL_PITCH_RANDOMIZATION_RANGE, TASK_SPEC::TARGET_FRAME_ROLL_PITCH_RANDOMIZATION_RANGE, rng);
+                parameters.target_pitch = random::uniform_real_distribution(device.random, -TASK_SPEC::TARGET_FRAME_ROLL_PITCH_RANDOMIZATION_RANGE, TASK_SPEC::TARGET_FRAME_ROLL_PITCH_RANDOMIZATION_RANGE, rng);
+            } else {
+                parameters.target_roll = 0;
+                parameters.target_pitch = 0;
+            }
+            if constexpr (TASK_SPEC::TARGET_FRAME_BRIGHTNESS_MISMATCH_RANGE > (T)0){
+                parameters.brightness_mismatch = (T)1 + (random::uniform_real_distribution(device.random, (T)0, (T)1, rng) * (T)2 - (T)1) * TASK_SPEC::TARGET_FRAME_BRIGHTNESS_MISMATCH_RANGE;
+            } else {
+                parameters.brightness_mismatch = 1;
+            }
+        }
+    }
+
+    template <typename DEVICE, typename TASK_SPEC, typename DATASET>
+    void init(DEVICE& device, rl::environments::hyperdrone::tasks::target_frame::World<TASK_SPEC>& world, typename TASK_SPEC::NEXT_WORLD::SharedContext& shared, const DATASET& dataset, const typename DATASET::Corpus& corpus, typename TASK_SPEC::TI first_scene, typename TASK_SPEC::TI num_scenes, typename TASK_SPEC::TI member_index) {
+        using NEXT_WORLD = typename TASK_SPEC::NEXT_WORLD;
+        init(device, static_cast<NEXT_WORLD&>(world), shared, dataset, corpus, first_scene, num_scenes, member_index);
+        malloc(device, world.target_frames);
+    }
+    template <typename DEVICE, typename TASK_SPEC>
+    void free(DEVICE& device, rl::environments::hyperdrone::tasks::target_frame::World<TASK_SPEC>& world) {
+        using NEXT_WORLD = typename TASK_SPEC::NEXT_WORLD;
+        if (!world.slots.empty()) {
+            free(device, world.target_frames);
+        }
+        free(device, static_cast<NEXT_WORLD&>(world));
+    }
+
+    template <typename DEVICE, typename TASK_SPEC, typename RNG>
+    void sample_initial_parameters(DEVICE& device, rl::environments::hyperdrone::tasks::target_frame::World<TASK_SPEC>& world, typename rl::environments::hyperdrone::tasks::target_frame::World<TASK_SPEC>::Parameters& parameters, RNG& rng) {
+        rl::environments::hyperdrone::tasks::target_frame::_sample_initial_parameters<DEVICE, TASK_SPEC>(device, world.dynamics, world.parameters, world.drone_parameters, world.num_drones, parameters, rng);
+    }
+    template <typename DEVICE, typename TASK_SPEC, typename PARAMETER_SPEC, typename RESET_SPEC, typename RNG>
+    void sample_initial_parameters(DEVICE& device, rl::environments::hyperdrone::tasks::target_frame::World<TASK_SPEC>& world, Tensor<PARAMETER_SPEC>& parameters, const Tensor<RESET_SPEC>& reset_mask, RNG& rng) {
+        using TI = typename TASK_SPEC::TI;
+        using WORLD = rl::environments::hyperdrone::tasks::target_frame::World<TASK_SPEC>;
+        static_assert(utils::typing::is_same_v<typename PARAMETER_SPEC::T, typename WORLD::Parameters>);
+        for (TI instance_i = 0; instance_i < WORLD::INSTANCES; instance_i++) {
+            if (get(device, reset_mask, instance_i)) {
+                sample_initial_parameters(device, world, get_ref(device, parameters, instance_i), rng);
+            }
+        }
+    }
+
+    // reset-conditional second render pass: the target frames are cached at reset (pixel-identical
+    // to per-step re-rendering since pose and brightness are fixed per episode), then the student
+    // frame render forwards to the base
+    template <typename DEVICE, typename TASK_SPEC, typename PARAMETER_SPEC, typename STATE_SPEC, typename RESET_SPEC>
+    void render(DEVICE& device, rl::environments::hyperdrone::tasks::target_frame::World<TASK_SPEC>& world, Tensor<PARAMETER_SPEC>& parameters, Tensor<STATE_SPEC>& states, const Tensor<RESET_SPEC>& reset_mask_input) {
+        request_render(device, world, reset_mask_input);
+        auto& reset_mask = world.render_reset;
+        using T = typename TASK_SPEC::T;
+        using TI = typename TASK_SPEC::TI;
+        using WORLD = rl::environments::hyperdrone::tasks::target_frame::World<TASK_SPEC>;
+        using NEXT_WORLD = typename TASK_SPEC::NEXT_WORLD;
+        using BASE_SPEC = typename NEXT_WORLD::SPEC;
+        constexpr TI INSTANCES = NEXT_WORLD::INSTANCES;
+        bool any_reset = false;
+        for (TI instance_i = 0; instance_i < INSTANCES; instance_i++) {
+            any_reset = any_reset || get(device, reset_mask, instance_i);
+        }
+        if (any_reset) {
+            auto& renderer = get_renderer(device, world);
+            const T aspect = static_cast<T>(BASE_SPEC::CAM_WIDTH) / static_cast<T>(BASE_SPEC::CAM_HEIGHT);
+            for (TI instance_i = 0; instance_i < INSTANCES; instance_i++) {
+                const auto& instance_parameters = get_ref(device, parameters, instance_i);
+                set(device, world.camera_staging_close, rl::environments::hyperdrone::make_target_camera<DEVICE, T>(device, instance_parameters.camera, aspect, instance_parameters.scene_translation, instance_parameters.scene_yaw_cos, instance_parameters.scene_yaw_sin, instance_parameters.target_roll, instance_parameters.target_pitch), instance_i);
+            }
+            copy(device, renderer.device, world.camera_staging_close, cameras(device, renderer));
+            if constexpr (BASE_SPEC::ENABLE_MOTION_BLUR) {
+                copy(device, renderer.device, world.camera_staging_close, cameras_open(device, renderer));
+            }
+            render(device, renderer);
+            constexpr TI CAM_PIXELS = BASE_SPEC::CAM_WIDTH * BASE_SPEC::CAM_HEIGHT;
+            std::vector<float> observation_staging(INSTANCES * CAM_PIXELS * 3);
+            {
+                Tensor<tensor::Specification<float, TI, typename decltype(renderer.observation)::SPEC::SHAPE>> observation_alias;
+                observation_alias._data = observation_staging.data();
+                copy(renderer.device, device, renderer.observation, observation_alias);
+            }
+            std::vector<float> target_staging(INSTANCES * NEXT_WORLD::FRAME_DIM);
+            {
+                Tensor<typename WORLD::TARGET_FRAMES_SPEC> target_alias;
+                target_alias._data = target_staging.data();
+                copy(device, device, world.target_frames, target_alias);
+                for (TI instance_i = 0; instance_i < INSTANCES; instance_i++) {
+                    if (!get(device, reset_mask, instance_i)) {
+                        continue;
+                    }
+                    const auto& instance_parameters = get_ref(device, parameters, instance_i);
+                    const float scale = (float)(instance_parameters.brightness_scale * instance_parameters.brightness_mismatch);
+                    for (TI pixel_i = 0; pixel_i < CAM_PIXELS; pixel_i++) {
+                        for (TI channel_i = 0; channel_i < 3; channel_i++) {
+                            float value = observation_staging[(instance_i * CAM_PIXELS + pixel_i) * 3 + channel_i] * scale;
+                            value = value < 0.0f ? 0.0f : (value > 1.0f ? 1.0f : value);
+                            target_staging[instance_i * NEXT_WORLD::FRAME_DIM + pixel_i * NEXT_WORLD::IMAGE_CHANNELS + channel_i] = value;
+                        }
+                    }
+                }
+                copy(device, device, target_alias, world.target_frames);
+            }
+        }
+        render(device, static_cast<NEXT_WORLD&>(world), parameters, states, reset_mask);
+    }
+
+    template <typename DEVICE, typename TASK_SPEC, typename PARAMETER_SPEC, typename STATE_SPEC, typename OBSERVATION_SPEC, typename RNG>
+    void observe(DEVICE& device, rl::environments::hyperdrone::tasks::target_frame::World<TASK_SPEC>& world, Tensor<PARAMETER_SPEC>& parameters, Tensor<STATE_SPEC>& states, typename TASK_SPEC::NEXT_WORLD::Observation observation_type, Tensor<OBSERVATION_SPEC>& observations, RNG& rng){
+        using NEXT_WORLD = typename TASK_SPEC::NEXT_WORLD;
+        if(world.render_pending){
+            render(device, world, parameters, states, world.render_reset);
+        }
+        observe(device, static_cast<NEXT_WORLD&>(world), parameters, states, observation_type, observations, rng);
+    }
+
+    // the composed observation: [stacked frames | target frame | padding]
+    template <typename DEVICE, typename TASK_SPEC, typename PARAMETER_SPEC, typename STATE_SPEC, typename OBSERVATION_SPEC, typename RNG>
+    void observe(DEVICE& device, rl::environments::hyperdrone::tasks::target_frame::World<TASK_SPEC>& world, Tensor<PARAMETER_SPEC>& parameters, Tensor<STATE_SPEC>& states, typename rl::environments::hyperdrone::tasks::target_frame::World<TASK_SPEC>::Observation, Tensor<OBSERVATION_SPEC>& observations, RNG& rng) {
+        using T = typename TASK_SPEC::T;
+        using TI = typename TASK_SPEC::TI;
+        using WORLD = rl::environments::hyperdrone::tasks::target_frame::World<TASK_SPEC>;
+        using NEXT_WORLD = typename TASK_SPEC::NEXT_WORLD;
+        using BASE_SPEC = typename NEXT_WORLD::SPEC;
+        constexpr TI INSTANCES = NEXT_WORLD::INSTANCES;
+        constexpr TI CAM_PIXELS = BASE_SPEC::CAM_WIDTH * BASE_SPEC::CAM_HEIGHT;
+        constexpr TI IMAGE_CHANNELS = NEXT_WORLD::IMAGE_CHANNELS;
+        constexpr TI STACK_N = TASK_SPEC::IMAGE_STACK_N;
+        constexpr TI STACK_STRIDE = TASK_SPEC::IMAGE_STACK_STRIDE;
+        constexpr TI TOTAL_CHANNELS = WORLD::OBSERVATION_CHANNELS;
+        if(world.render_pending){
+            render(device, world, parameters, states, world.render_reset);
+        }
+        static_assert(get<0>(typename OBSERVATION_SPEC::SHAPE{}) == INSTANCES);
+        static_assert(get<1>(typename OBSERVATION_SPEC::SHAPE{}) == WORLD::OBSERVATION_DIM);
+        utils::assert_exit(device, world.history_step > 0, "hyperdrone::tasks::target_frame::observe: no frame available");
+        const TI latest = world.history_step - 1;
+
+        std::vector<TI> episode_start_staging(INSTANCES);
+        {
+            Tensor<typename NEXT_WORLD::EPISODE_START_SPEC> episode_start_alias;
+            episode_start_alias._data = episode_start_staging.data();
+            copy(device, device, world.episode_start, episode_start_alias);
+        }
+        std::vector<float> target_staging(INSTANCES * NEXT_WORLD::FRAME_DIM);
+        {
+            Tensor<typename WORLD::TARGET_FRAMES_SPEC> target_alias;
+            target_alias._data = target_staging.data();
+            copy(device, device, world.target_frames, target_alias);
+        }
+        std::vector<float> frame_staging(NEXT_WORLD::FRAME_DIM);
+        for (TI instance_i = 0; instance_i < INSTANCES; instance_i++) {
+            for (TI frame_i = 0; frame_i < STACK_N; frame_i++) {
+                const TI back = frame_i * STACK_STRIDE;
+                TI desired = latest >= back ? latest - back : episode_start_staging[instance_i];
+                if (desired < episode_start_staging[instance_i]) {
+                    desired = episode_start_staging[instance_i];
+                }
+                const TI history_slot = desired % BASE_SPEC::HISTORY_LENGTH;
+                {
+                    auto history_row = view(device, world.history, history_slot);
+                    auto history_instance = view(device, history_row, instance_i);
+                    Tensor<tensor::Specification<float, TI, tensor::Shape<TI, NEXT_WORLD::FRAME_DIM>>> frame_alias;
+                    frame_alias._data = frame_staging.data();
+                    copy(device, device, history_instance, frame_alias);
+                }
+                for (TI pixel_i = 0; pixel_i < CAM_PIXELS; pixel_i++) {
+                    for (TI channel_i = 0; channel_i < IMAGE_CHANNELS; channel_i++) {
+                        set(device, observations, (T)frame_staging[pixel_i * IMAGE_CHANNELS + channel_i], instance_i, pixel_i * TOTAL_CHANNELS + frame_i * IMAGE_CHANNELS + channel_i);
+                    }
+                }
+            }
+            for (TI pixel_i = 0; pixel_i < CAM_PIXELS; pixel_i++) {
+                for (TI channel_i = 0; channel_i < IMAGE_CHANNELS; channel_i++) {
+                    set(device, observations, (T)target_staging[instance_i * NEXT_WORLD::FRAME_DIM + pixel_i * IMAGE_CHANNELS + channel_i], instance_i, pixel_i * TOTAL_CHANNELS + STACK_N * IMAGE_CHANNELS + channel_i);
+                }
+                for (TI channel_i = (STACK_N + 1) * IMAGE_CHANNELS; channel_i < TOTAL_CHANNELS; channel_i++) {
+                    set(device, observations, (T)0, instance_i, pixel_i * TOTAL_CHANNELS + channel_i);
+                }
+            }
+        }
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+// ADL entry points: generic composite code (e.g. the MultiEnvironment lifecycle fan-out)
+// dispatches member verbs without Tensor arguments, so the task's lifecycle overloads must be
+// reachable through the member type's namespace
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::rl::environments::hyperdrone::tasks::target_frame {
+    template <typename DEVICE, typename TASK_SPEC, typename DATASET>
+    void init(DEVICE& device, World<TASK_SPEC>& world, typename TASK_SPEC::NEXT_WORLD::SharedContext& shared, const DATASET& dataset, const typename DATASET::Corpus& corpus, typename TASK_SPEC::TI first_scene, typename TASK_SPEC::TI num_scenes, typename TASK_SPEC::TI member_index){
+        ::rl_tools::init(device, world, shared, dataset, corpus, first_scene, num_scenes, member_index);
+    }
+    using rl_tools::free;
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/environments/hyperdrone/tasks/target_frame/operations_cuda.h

@@ -0,0 +1,218 @@
+#include "../../../../../version.h"
+#include "../../../../../rl_tools.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ENVIRONMENTS_HYPERDRONE_TASKS_TARGET_FRAME_OPERATIONS_CUDA_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ENVIRONMENTS_HYPERDRONE_TASKS_TARGET_FRAME_OPERATIONS_CUDA_H
+
+#include "operations_cpu.h"
+#include "../../operations_cuda.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template <typename DEV_SPEC, typename TASK_SPEC>
+    void free(devices::CUDA<DEV_SPEC>& device, rl::environments::hyperdrone::tasks::target_frame::World<TASK_SPEC>& world){
+        free<devices::CUDA<DEV_SPEC>, TASK_SPEC>(device, world);
+        if(world.cuda_reset_staging != nullptr){
+            cudaFreeHost(world.cuda_reset_staging);
+            world.cuda_reset_staging = nullptr;
+        }
+        check_status(device);
+    }
+
+    namespace rl::environments::hyperdrone::tasks::target_frame::cuda{
+        template <typename DEVICE, typename TASK_SPEC, typename DRONE_SPEC, typename PARAMETER_SPEC, typename RESET_SPEC, typename RNG>
+        __global__ void sample_initial_parameters_kernel(DEVICE device, typename World<TASK_SPEC>::NEXT_WORLD::DYNAMICS_ENV dynamics, typename World<TASK_SPEC>::NEXT_WORLD::Parameters defaults, Tensor<DRONE_SPEC> drones, typename TASK_SPEC::TI num_drones, Tensor<PARAMETER_SPEC> parameters, const Tensor<RESET_SPEC> reset_mask, RNG rng){
+            using TI = typename DEVICE::index_t;
+            constexpr TI INSTANCES = World<TASK_SPEC>::INSTANCES;
+            static_assert(RNG::NUM_RNGS >= INSTANCES, "Please increase the number of CUDA RNGs");
+            TI instance_i = threadIdx.x + blockIdx.x * blockDim.x;
+            if(instance_i < INSTANCES && get(device, reset_mask, instance_i)){
+                auto& rng_state = get(rng.states, 0, instance_i);
+                // qualified: ADL through the base-typed defaults argument would otherwise also
+                // consider the base helper and hard-instantiate the base World with TASK_SPEC
+                tasks::target_frame::_sample_initial_parameters<DEVICE, TASK_SPEC>(device, dynamics, defaults, drones, num_drones, get_ref(device, parameters, instance_i), rng_state);
+            }
+        }
+        template <typename DEVICE, typename TASK_SPEC, typename PARAMETER_SPEC>
+        __global__ void target_camera_kernel(DEVICE device, Tensor<PARAMETER_SPEC> parameters, rendering::raytracing::Camera<typename TASK_SPEC::T>* cameras, rendering::raytracing::Camera<typename TASK_SPEC::T>* cameras_open, typename TASK_SPEC::T aspect){
+            using T = typename TASK_SPEC::T;
+            using TI = typename DEVICE::index_t;
+            constexpr TI INSTANCES = World<TASK_SPEC>::INSTANCES;
+            TI instance_i = threadIdx.x + blockIdx.x * blockDim.x;
+            if(instance_i < INSTANCES){
+                const auto& instance_parameters = get_ref(device, parameters, instance_i);
+                auto camera = make_target_camera<DEVICE, T>(device, instance_parameters.camera, aspect, instance_parameters.scene_translation, instance_parameters.scene_yaw_cos, instance_parameters.scene_yaw_sin, instance_parameters.target_roll, instance_parameters.target_pitch);
+                cameras[instance_i] = camera;
+                if(cameras_open != cameras){
+                    cameras_open[instance_i] = camera;
+                }
+            }
+        }
+        template <typename DEVICE, typename TASK_SPEC, typename PARAMETER_SPEC, typename RESET_SPEC>
+        __global__ void target_scatter_kernel(DEVICE device, const float* __restrict__ observation, float* __restrict__ target_frames, Tensor<PARAMETER_SPEC> parameters, const Tensor<RESET_SPEC> reset_mask){
+            using TI = typename DEVICE::index_t;
+            using WORLD = World<TASK_SPEC>;
+            using NEXT_WORLD = typename TASK_SPEC::NEXT_WORLD;
+            using BASE_SPEC = typename NEXT_WORLD::SPEC;
+            constexpr TI INSTANCES = WORLD::INSTANCES;
+            constexpr TI CAM_PIXELS = BASE_SPEC::CAM_WIDTH * BASE_SPEC::CAM_HEIGHT;
+            TI thread_i = threadIdx.x + blockIdx.x * blockDim.x;
+            if(thread_i < INSTANCES * CAM_PIXELS){
+                TI instance_i = thread_i / CAM_PIXELS;
+                if(!get(device, reset_mask, instance_i)){
+                    return;
+                }
+                TI pixel_i = thread_i % CAM_PIXELS;
+                const auto& instance_parameters = get_ref(device, parameters, instance_i);
+                const float scale = (float)(instance_parameters.brightness_scale * instance_parameters.brightness_mismatch);
+                for(TI channel_i = 0; channel_i < 3; channel_i++){
+                    float value = observation[(instance_i * CAM_PIXELS + pixel_i) * 3 + channel_i] * scale;
+                    value = value < 0.0f ? 0.0f : (value > 1.0f ? 1.0f : value);
+                    target_frames[instance_i * NEXT_WORLD::FRAME_DIM + pixel_i * NEXT_WORLD::IMAGE_CHANNELS + channel_i] = value;
+                }
+            }
+        }
+        template <typename DEVICE, typename TASK_SPEC, typename OBSERVATION_SPEC>
+        __global__ void observe_kernel(DEVICE device, const float* __restrict__ history, const typename TASK_SPEC::TI* __restrict__ episode_start, const float* __restrict__ target_frames, typename TASK_SPEC::TI latest, Tensor<OBSERVATION_SPEC> observations){
+            using T = typename TASK_SPEC::T;
+            using TI = typename DEVICE::index_t;
+            using WORLD = World<TASK_SPEC>;
+            using NEXT_WORLD = typename TASK_SPEC::NEXT_WORLD;
+            using BASE_SPEC = typename NEXT_WORLD::SPEC;
+            constexpr TI INSTANCES = WORLD::INSTANCES;
+            constexpr TI CAM_PIXELS = BASE_SPEC::CAM_WIDTH * BASE_SPEC::CAM_HEIGHT;
+            constexpr TI IMAGE_CHANNELS = NEXT_WORLD::IMAGE_CHANNELS;
+            constexpr TI TOTAL_CHANNELS = WORLD::OBSERVATION_CHANNELS;
+            constexpr TI STACK_N = TASK_SPEC::IMAGE_STACK_N;
+            constexpr TI STACK_STRIDE = TASK_SPEC::IMAGE_STACK_STRIDE;
+            TI thread_i = threadIdx.x + blockIdx.x * blockDim.x;
+            if(thread_i < INSTANCES * WORLD::OBSERVATION_DIM){
+                TI instance_i = thread_i / WORLD::OBSERVATION_DIM;
+                TI offset = thread_i % WORLD::OBSERVATION_DIM;
+                TI pixel_i = offset / TOTAL_CHANNELS;
+                TI channel = offset % TOTAL_CHANNELS;
+                float value;
+                if(channel < STACK_N * IMAGE_CHANNELS){
+                    TI frame_i = channel / IMAGE_CHANNELS;
+                    TI channel_i = channel % IMAGE_CHANNELS;
+                    TI back = frame_i * STACK_STRIDE;
+                    TI start = episode_start[instance_i];
+                    TI desired = latest >= back ? latest - back : start;
+                    if(desired < start){
+                        desired = start;
+                    }
+                    TI history_slot = desired % BASE_SPEC::HISTORY_LENGTH;
+                    value = history[(history_slot * INSTANCES * NEXT_WORLD::N_VIEWS + instance_i) * NEXT_WORLD::FRAME_DIM + pixel_i * IMAGE_CHANNELS + channel_i];
+                }
+                else if(channel < (STACK_N + 1) * IMAGE_CHANNELS){
+                    TI channel_i = channel - STACK_N * IMAGE_CHANNELS;
+                    value = target_frames[instance_i * NEXT_WORLD::FRAME_DIM + pixel_i * IMAGE_CHANNELS + channel_i];
+                }
+                else{
+                    value = 0.0f;
+                }
+                set(device, observations, (T)value, instance_i, offset);
+            }
+        }
+    }
+
+    template <typename DEV_SPEC, typename TASK_SPEC, typename PARAMETER_SPEC, typename RESET_SPEC, typename RNG>
+    void sample_initial_parameters(devices::CUDA<DEV_SPEC>& device, rl::environments::hyperdrone::tasks::target_frame::World<TASK_SPEC>& world, Tensor<PARAMETER_SPEC>& parameters, const Tensor<RESET_SPEC>& reset_mask, RNG& rng){
+        using DEVICE = devices::CUDA<DEV_SPEC>;
+        using TI = typename DEVICE::index_t;
+        constexpr TI INSTANCES = rl::environments::hyperdrone::tasks::target_frame::World<TASK_SPEC>::INSTANCES;
+        constexpr TI BLOCKSIZE = 32;
+        constexpr TI N_BLOCKS = RL_TOOLS_DEVICES_CUDA_CEIL(INSTANCES, BLOCKSIZE);
+        devices::cuda::TAG<DEVICE, true> tag_device{};
+        auto world_rng = rl::environments::hyperdrone::cuda::instance_rng(device, world, rng);
+        rl::environments::hyperdrone::tasks::target_frame::cuda::sample_initial_parameters_kernel<decltype(tag_device), TASK_SPEC><<<dim3(N_BLOCKS), dim3(BLOCKSIZE), 0, device.stream>>>(tag_device, world.dynamics, world.parameters, world.drone_parameters, world.num_drones, parameters, reset_mask, world_rng);
+        check_status(device);
+    }
+    template <typename DEV_SPEC, typename TASK_SPEC, typename PARAMETER_SPEC, typename STATE_SPEC, typename RESET_SPEC>
+    void render(devices::CUDA<DEV_SPEC>& device, rl::environments::hyperdrone::tasks::target_frame::World<TASK_SPEC>& world, Tensor<PARAMETER_SPEC>& parameters, Tensor<STATE_SPEC>& states, const Tensor<RESET_SPEC>& reset_mask_input){
+        request_render(device, world, reset_mask_input);
+        auto& reset_mask = world.render_reset;
+        using DEVICE = devices::CUDA<DEV_SPEC>;
+        using T = typename TASK_SPEC::T;
+        using TI = typename DEVICE::index_t;
+        using WORLD = rl::environments::hyperdrone::tasks::target_frame::World<TASK_SPEC>;
+        using NEXT_WORLD = typename TASK_SPEC::NEXT_WORLD;
+        using BASE_SPEC = typename NEXT_WORLD::SPEC;
+        constexpr TI INSTANCES = NEXT_WORLD::INSTANCES;
+        constexpr TI CAM_PIXELS = BASE_SPEC::CAM_WIDTH * BASE_SPEC::CAM_HEIGHT;
+        // the cache-refresh decision needs the reset mask on the host: a pinned mirror and a
+        // short join on the producing stream (the payoff is skipping a whole render pass on
+        // non-reset steps)
+        if(world.cuda_reset_staging == nullptr){
+            cudaMallocHost(&world.cuda_reset_staging, INSTANCES * sizeof(bool));
+        }
+        bool* reset_staging = (bool*)world.cuda_reset_staging;
+        cudaMemcpyAsync(reset_staging, data(reset_mask), INSTANCES * sizeof(bool), cudaMemcpyDeviceToHost, device.stream);
+        cudaStreamSynchronize(device.stream);
+        bool any_reset = false;
+        for(TI instance_i = 0; instance_i < INSTANCES; instance_i++){
+            any_reset = any_reset || reset_staging[instance_i];
+        }
+        if(any_reset){
+            auto& renderer = get_renderer(device, world);
+            cudaStream_t render_stream = stream(device, renderer);
+            rl::environments::hyperdrone::cuda::stream_barrier(world, device.stream, render_stream);
+            devices::cuda::TAG<DEVICE, true> tag_device{};
+            const T aspect = static_cast<T>(BASE_SPEC::CAM_WIDTH) / static_cast<T>(BASE_SPEC::CAM_HEIGHT);
+            {
+                constexpr TI BLOCKSIZE = 32;
+                constexpr TI N_BLOCKS = RL_TOOLS_DEVICES_CUDA_CEIL(INSTANCES, BLOCKSIZE);
+                auto* cameras_data = data(cameras(device, renderer));
+                auto* cameras_open_data = cameras_data;
+                if constexpr(BASE_SPEC::ENABLE_MOTION_BLUR){
+                    cameras_open_data = data(cameras_open(device, renderer));
+                }
+                rl::environments::hyperdrone::tasks::target_frame::cuda::target_camera_kernel<decltype(tag_device), TASK_SPEC><<<dim3(N_BLOCKS), dim3(BLOCKSIZE), 0, render_stream>>>(tag_device, parameters, cameras_data, cameras_open_data, aspect);
+            }
+            render_launch(device, renderer);
+            {
+                constexpr TI BLOCKSIZE = 256;
+                constexpr TI N_BLOCKS = RL_TOOLS_DEVICES_CUDA_CEIL(INSTANCES * CAM_PIXELS, BLOCKSIZE);
+                rl::environments::hyperdrone::tasks::target_frame::cuda::target_scatter_kernel<decltype(tag_device), TASK_SPEC><<<dim3(N_BLOCKS), dim3(BLOCKSIZE), 0, render_stream>>>(tag_device, data(renderer.observation), data(world.target_frames), parameters, reset_mask);
+            }
+            check_status(device);
+        }
+        render(device, static_cast<NEXT_WORLD&>(world), parameters, states, reset_mask);
+    }
+    template <typename DEV_SPEC, typename TASK_SPEC, typename PARAMETER_SPEC, typename STATE_SPEC, typename OBSERVATION_SPEC, typename RNG>
+    void observe(devices::CUDA<DEV_SPEC>& device, rl::environments::hyperdrone::tasks::target_frame::World<TASK_SPEC>& world, Tensor<PARAMETER_SPEC>& parameters, Tensor<STATE_SPEC>& states, typename TASK_SPEC::NEXT_WORLD::Observation observation_type, Tensor<OBSERVATION_SPEC>& observations, RNG& rng){
+        using NEXT_WORLD = typename TASK_SPEC::NEXT_WORLD;
+        if(world.render_pending){
+            render(device, world, parameters, states, world.render_reset);
+        }
+        observe(device, static_cast<NEXT_WORLD&>(world), parameters, states, observation_type, observations, rng);
+    }
+    template <typename DEV_SPEC, typename TASK_SPEC, typename PARAMETER_SPEC, typename STATE_SPEC, typename OBSERVATION_SPEC, typename RNG>
+    void observe(devices::CUDA<DEV_SPEC>& device, rl::environments::hyperdrone::tasks::target_frame::World<TASK_SPEC>& world, Tensor<PARAMETER_SPEC>& parameters, Tensor<STATE_SPEC>& states, typename rl::environments::hyperdrone::tasks::target_frame::World<TASK_SPEC>::Observation, Tensor<OBSERVATION_SPEC>& observations, RNG& rng){
+        using DEVICE = devices::CUDA<DEV_SPEC>;
+        using TI = typename DEVICE::index_t;
+        using WORLD = rl::environments::hyperdrone::tasks::target_frame::World<TASK_SPEC>;
+        static_assert(get<0>(typename OBSERVATION_SPEC::SHAPE{}) == WORLD::INSTANCES);
+        static_assert(get<1>(typename OBSERVATION_SPEC::SHAPE{}) == WORLD::OBSERVATION_DIM);
+        if(world.render_pending){
+            render(device, world, parameters, states, world.render_reset);
+        }
+        utils::assert_exit(device, world.history_step > 0, "hyperdrone::tasks::target_frame::observe: no frame available");
+        cudaStream_t render_stream = stream(device, get_renderer(device, world));
+        devices::cuda::TAG<DEVICE, true> tag_device{};
+        constexpr TI BLOCKSIZE = 256;
+        constexpr TI N_BLOCKS = RL_TOOLS_DEVICES_CUDA_CEIL(WORLD::INSTANCES * WORLD::OBSERVATION_DIM, BLOCKSIZE);
+        rl::environments::hyperdrone::tasks::target_frame::cuda::observe_kernel<decltype(tag_device), TASK_SPEC><<<dim3(N_BLOCKS), dim3(BLOCKSIZE), 0, render_stream>>>(tag_device, data(world.history), data(world.episode_start), data(world.target_frames), world.history_step - 1, observations);
+        rl::environments::hyperdrone::cuda::stream_barrier(world, render_stream, device.stream);
+        check_status(device);
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::rl::environments::hyperdrone::tasks::target_frame {
+    using rl_tools::free;
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/environments/hyperdrone/tasks/target_frame/operations_generic_position_hold.h

@@ -0,0 +1,128 @@
+#include "../../../../../version.h"
+#include "../../../../../rl_tools.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ENVIRONMENTS_HYPERDRONE_TASKS_TARGET_FRAME_OPERATIONS_GENERIC_POSITION_HOLD_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ENVIRONMENTS_HYPERDRONE_TASKS_TARGET_FRAME_OPERATIONS_GENERIC_POSITION_HOLD_H
+
+#include "position_hold.h"
+#include "../../../../../nn_models/sequential/operations_generic.h"
+#include "../../../../../nn_models/parallel/operations_generic.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template<typename DEVICE, typename MODEL, typename TYPE_POLICY, typename TI>
+    RL_TOOLS_FUNCTION_PLACEMENT void init(DEVICE& device, MODEL& actor, const rl::environments::hyperdrone::tasks::target_frame::position_hold::Specification<TYPE_POLICY, TI>&){
+        using SPEC = rl::environments::hyperdrone::tasks::target_frame::position_hold::Specification<TYPE_POLICY, TI>;
+        using T = typename SPEC::T;
+        constexpr TI ACTION_DIM = SPEC::ACTION_DIM;
+        using nn_models::sequential::layer;
+        {
+            auto& image_branch = get<0>(actor.pipelines);
+            auto& grey = layer<0>(image_branch);
+            auto& cost = layer<1>(image_branch);
+            auto& pool = layer<2>(image_branch);
+            set_all(device, layer<5>(image_branch).biases.parameters, 0);
+            set_all(device, layer<6>(image_branch).biases.parameters, 0);
+            set_all(device, layer<7>(image_branch).biases.parameters, 0);
+            set_all(device, layer<7>(image_branch).weights.parameters, 0);
+            auto& pair_difference = layer<9>(image_branch);
+            set_all(device, pair_difference.weights.parameters, 0);
+            set_all(device, pair_difference.biases.parameters, 0);
+            for(TI action = 0; action < ACTION_DIM; action++){
+                set(device, pair_difference.weights.parameters, static_cast<T>(0.5), action, action);
+                set(device, pair_difference.weights.parameters, static_cast<T>(-0.5), action, action + ACTION_DIM);
+            }
+            auto& known = layer<1>(get<1>(actor.pipelines));
+            using STATE_WEIGHTS_SHAPE = typename decltype(known.weights.parameters)::SPEC::SHAPE;
+            static_assert(get<0>(STATE_WEIGHTS_SHAPE{}) == 4 && get<1>(STATE_WEIGHTS_SHAPE{}) == 65);
+            set_all(device, known.weights.parameters, 0);
+            set_all(device, known.biases.parameters, 0);
+            const T gravity = SPEC::GRAVITY;
+            set(device, known.weights.parameters, static_cast<T>(1), 0, 6);
+            set(device, known.weights.parameters, static_cast<T>(-0.02) * gravity, 0, 40);
+            set(device, known.weights.parameters, static_cast<T>(-0.04), 0, 64);
+            set(device, known.weights.parameters, static_cast<T>(-1.5), 1, 1);
+            set(device, known.weights.parameters, static_cast<T>(3) / gravity, 1, 63);
+            set(device, known.weights.parameters, static_cast<T>(1.5), 2, 0);
+            set(device, known.weights.parameters, static_cast<T>(-3) / gravity, 2, 62);
+            set_all(device, actor.head.input_layer.weights.parameters, 0);
+            set_all(device, actor.head.input_layer.biases.parameters, 0);
+            set_all(device, actor.head.output_layer.weights.parameters, 0);
+            set_all(device, actor.head.output_layer.biases.parameters, 0);
+            for(TI action = 0; action < ACTION_DIM; action++){
+                set(device, actor.head.input_layer.weights.parameters, static_cast<T>(1), action, action);
+                set(device, actor.head.input_layer.weights.parameters, static_cast<T>(1), action, ACTION_DIM + action);
+                set(device, actor.head.output_layer.weights.parameters, static_cast<T>(1), action, action);
+            }
+            set_all(device, grey.weights.parameters, 0);
+            set_all(device, grey.biases.parameters, 0);
+            set_all(device, cost.weights.parameters, 0);
+            set_all(device, cost.biases.parameters, 0);
+            set_all(device, pool.weights.parameters, 0);
+            set_all(device, pool.biases.parameters, 0);
+            auto& normalization = layer<3>(image_branch);
+            auto& difference = layer<4>(image_branch);
+            set_all(device, normalization.weights.parameters, 0);
+            set_all(device, normalization.biases.parameters, 0);
+            set_all(device, difference.weights.parameters, 0);
+            set_all(device, difference.biases.parameters, 0);
+            const T scales[5] = {0.5, 2, 8, 32, 128};
+            for(TI feature = 0; feature < 25; feature++){
+                for(TI level = 0; level < 5; level++){
+                    for(TI direction = 0; direction < 2; direction++){
+                        const TI channel = (feature * 5 + level) * 2 + direction;
+                        set(device, normalization.weights.parameters, -scales[level], channel, 0, 0, feature * 2 + direction);
+                        set(device, difference.weights.parameters, static_cast<T>(1), feature * 2 + direction, 0, 0, channel);
+                    }
+                }
+            }
+            const TI frame_channels[2] = {0, 3};
+            const T luminance[3] = {0.299, 0.587, 0.114};
+            for(TI frame = 0; frame < 2; frame++){
+                for(TI y = 0; y < 2; y++){
+                    for(TI x = 0; x < 2; x++){
+                        for(TI channel = 0; channel < 3; channel++){
+                            set(device, grey.weights.parameters, luminance[channel] / static_cast<T>(4), frame, y, x, frame_channels[frame] + channel);
+                        }
+                    }
+                }
+            }
+            for(TI frame = 0; frame < 1; frame++){
+                for(TI y = 0; y < 5; y++){
+                    for(TI x = 0; x < 5; x++){
+                        const TI feature = frame * 25 + y * 5 + x;
+                        for(TI direction = 0; direction < 2; direction++){
+                            const TI current = direction == 0 ? 0 : frame + 1;
+                            const TI reference = direction == 0 ? frame + 1 : 0;
+                            for(TI sign = 0; sign < 2; sign++){
+                                const TI output = feature * 4 + direction * 2 + sign;
+                                const T scale = sign == 0 ? static_cast<T>(10) : static_cast<T>(-10);
+                                set(device, cost.weights.parameters, scale, output, 2, 2, current);
+                                set(device, cost.weights.parameters, -scale, output, y, x, reference);
+                                for(TI py = 0; py < 2; py++){
+                                    for(TI px = 0; px < 2; px++){
+                                        const T interval_scale = 1;
+                                        set(device, pool.weights.parameters, static_cast<T>(1) / 4 * interval_scale, feature * 2 + direction, py, px, output);
+                                    }
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+        }
+        {
+            auto& state_normalization = layer<0>(get<1>(actor.pipelines));
+            const T gravity = SPEC::GRAVITY;
+            constexpr TI acceleration_offset = SPEC::BASE_OBSERVATION_DIM - 3 * SPEC::ACTION_HISTORY_LENGTH;
+            for(TI item = 0; item < SPEC::ACTION_HISTORY_LENGTH; item++){
+                for(TI axis = 0; axis < 3; axis++){
+                    set(device, state_normalization.mean.parameters, axis == 2 ? gravity : static_cast<T>(0), acceleration_offset + 3 * item + axis);
+                    set(device, state_normalization.precision.parameters, static_cast<T>(1) / gravity, acceleration_offset + 3 * item + axis);
+                }
+            }
+        }
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/environments/hyperdrone/tasks/target_frame/position_hold.h

@@ -0,0 +1,22 @@
+#include "../../../../../version.h"
+#include "../../../../../rl_tools.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ENVIRONMENTS_HYPERDRONE_TASKS_TARGET_FRAME_POSITION_HOLD_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ENVIRONMENTS_HYPERDRONE_TASKS_TARGET_FRAME_POSITION_HOLD_H
+
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::rl::environments::hyperdrone::tasks::target_frame::position_hold{
+    template<typename T_TYPE_POLICY, typename T_TI>
+    struct Specification{
+        using TYPE_POLICY = T_TYPE_POLICY;
+        using T = typename TYPE_POLICY::DEFAULT;
+        using TI = T_TI;
+        static constexpr T GRAVITY = 9.81;
+        static constexpr TI ACTION_HISTORY_LENGTH = 8;
+        static constexpr TI BASE_OBSERVATION_DIM = 62;
+        static constexpr TI ACTION_DIM = 4;
+    };
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif