hyperray 0.0.1__py3-none-any.whl

hyperdrone/_vendor/rl-tools/include/rl_tools/nn/optimizers/sgd/instance/operations_generic.h

@@ -0,0 +1,140 @@
+#include "../../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_OPTIMIZERS_SGD_INSTANCE_OPERATIONS_GENERIC_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_OPTIMIZERS_SGD_INSTANCE_OPERATIONS_GENERIC_H
+
+#include "../sgd.h"
+#include "../../../../nn/parameters/operations_generic.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::parameters::SGD::Instance<SPEC>& p){
+        malloc(device, (nn::parameters::Gradient::Instance<SPEC>&) p);
+        malloc(device, p.velocity);
+        if constexpr(nn::parameters::SGD::Instance<SPEC>::USE_MASTER_PARAMETERS){
+            malloc(device, p.master_parameters);
+        }
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::parameters::SGD::Instance<SPEC>& p){
+        free(device, (nn::parameters::Gradient::Instance<SPEC>&) p);
+        free(device, p.velocity);
+        if constexpr(nn::parameters::SGD::Instance<SPEC>::USE_MASTER_PARAMETERS){
+            free(device, p.master_parameters);
+        }
+    }
+    template<typename DEVICE, typename SPEC, typename PARAMETER_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void update(DEVICE& device, nn::parameters::SGD::Instance<PARAMETER_SPEC>& parameter, nn::optimizers::SGD<SPEC>& optimizer){
+        // SGD with momentum: v = momentum * v + gradient (+ weight_decay * param)
+        // Nesterov: param -= lr * (momentum * v + gradient)
+        // Classical: param -= lr * v
+        using TI = typename DEVICE::index_t;
+        using PARAMETER_INSTANCE = nn::parameters::SGD::Instance<PARAMETER_SPEC>;
+        static constexpr bool USE_MASTER_PARAMETERS = PARAMETER_INSTANCE::USE_MASTER_PARAMETERS;
+        using T_VELOCITY = typename PARAMETER_INSTANCE::T_VELOCITY;
+        using T_PARAMETER = typename decltype(parameter.parameters)::T;
+        using T_MASTER_PARAMETER = typename PARAMETER_INSTANCE::T_MASTER_PARAMETER;
+        const auto& optimizer_parameters = get(device, optimizer.parameters, 0);
+        auto params = matrix_view(device, parameter.parameters);
+        auto grad = matrix_view(device, parameter.gradient);
+        auto vel = matrix_view(device, parameter.velocity);
+        if constexpr(USE_MASTER_PARAMETERS){
+            auto master_params = matrix_view(device, parameter.master_parameters);
+            constexpr TI ROWS = decltype(params)::ROWS;
+            constexpr TI COLS = decltype(params)::COLS;
+            for(TI row_i = 0; row_i < ROWS; row_i++){
+                for(TI col_i = 0; col_i < COLS; col_i++){
+                    T_VELOCITY value = (T_VELOCITY)get(master_params, row_i, col_i);
+                    T_VELOCITY g = get(grad, row_i, col_i);
+                    if constexpr(SPEC::ENABLE_WEIGHT_DECAY){
+                        if constexpr(utils::typing::is_same_v<typename PARAMETER_SPEC::CATEGORY_TAG, nn::parameters::categories::Weights>){
+                            g += value * optimizer_parameters.weight_decay;
+                        }
+                    }
+                    T_VELOCITY v = optimizer_parameters.momentum * get(vel, row_i, col_i) + g;
+                    set(vel, row_i, col_i, v);
+                    T_VELOCITY param_update = optimizer_parameters.nesterov ? optimizer_parameters.momentum * v + g : v;
+                    value -= optimizer_parameters.learning_rate * param_update;
+                    set(master_params, row_i, col_i, (T_MASTER_PARAMETER)value);
+                    set(params, row_i, col_i, (T_PARAMETER)value);
+                }
+            }
+        }
+        else{
+            constexpr TI ROWS = decltype(params)::ROWS;
+            constexpr TI COLS = decltype(params)::COLS;
+            for(TI row_i = 0; row_i < ROWS; row_i++){
+                for(TI col_i = 0; col_i < COLS; col_i++){
+                    T_VELOCITY value = (T_VELOCITY)get(params, row_i, col_i);
+                    T_VELOCITY g = get(grad, row_i, col_i);
+                    if constexpr(SPEC::ENABLE_WEIGHT_DECAY){
+                        if constexpr(utils::typing::is_same_v<typename PARAMETER_SPEC::CATEGORY_TAG, nn::parameters::categories::Weights>){
+                            g += value * optimizer_parameters.weight_decay;
+                        }
+                    }
+                    T_VELOCITY v = optimizer_parameters.momentum * get(vel, row_i, col_i) + g;
+                    set(vel, row_i, col_i, v);
+                    T_VELOCITY param_update = optimizer_parameters.nesterov ? optimizer_parameters.momentum * v + g : v;
+                    value -= optimizer_parameters.learning_rate * param_update;
+                    set(params, row_i, col_i, (T_PARAMETER)value);
+                }
+            }
+        }
+    }
+    template<typename DEVICE, typename SPEC, typename PARAMETER_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void _reset_optimizer_state(DEVICE& device, nn::parameters::SGD::Instance<PARAMETER_SPEC>& parameter, nn::optimizers::SGD<SPEC>& optimizer){
+        set_all(device, parameter.velocity, 0);
+        if constexpr(nn::parameters::SGD::Instance<PARAMETER_SPEC>::USE_MASTER_PARAMETERS){
+            copy(device, device, parameter.parameters, parameter.master_parameters);
+        }
+    }
+    template<typename SOURCE_DEVICE, typename TARGET_DEVICE, typename SOURCE_SPEC, typename TARGET_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void copy(SOURCE_DEVICE& source_device, TARGET_DEVICE& target_device, const nn::parameters::SGD::Instance<SOURCE_SPEC>& source, nn::parameters::SGD::Instance<TARGET_SPEC>& target){
+        copy(source_device, target_device, (const nn::parameters::Gradient::Instance<SOURCE_SPEC>&) source, (nn::parameters::Gradient::Instance<TARGET_SPEC>&) target);
+        copy(source_device, target_device, source.velocity, target.velocity);
+        if constexpr(nn::parameters::SGD::Instance<SOURCE_SPEC>::USE_MASTER_PARAMETERS && nn::parameters::SGD::Instance<TARGET_SPEC>::USE_MASTER_PARAMETERS){
+            copy(source_device, target_device, source.master_parameters, target.master_parameters);
+        }
+        else if constexpr(!nn::parameters::SGD::Instance<SOURCE_SPEC>::USE_MASTER_PARAMETERS && nn::parameters::SGD::Instance<TARGET_SPEC>::USE_MASTER_PARAMETERS){
+            copy(source_device, target_device, source.parameters, target.master_parameters);
+        }
+    }
+    template<typename DEVICE, typename SPEC_1, typename SPEC_2>
+    RL_TOOLS_FUNCTION_PLACEMENT typename SPEC_1::TYPE_POLICY::DEFAULT abs_diff(DEVICE& device, const nn::parameters::SGD::Instance<SPEC_1>& p1, const nn::parameters::SGD::Instance<SPEC_2>& p2){
+        typename SPEC_1::TYPE_POLICY::DEFAULT acc = 0;
+        acc += abs_diff(device, static_cast<const nn::parameters::Gradient::Instance<SPEC_1>&>(p1), static_cast<const nn::parameters::Gradient::Instance<SPEC_2>&>(p2));
+        acc += abs_diff(device, p1.velocity, p2.velocity);
+        if constexpr(nn::parameters::SGD::Instance<SPEC_1>::USE_MASTER_PARAMETERS && nn::parameters::SGD::Instance<SPEC_2>::USE_MASTER_PARAMETERS){
+            acc += abs_diff(device, p1.master_parameters, p2.master_parameters);
+        }
+        else if constexpr(nn::parameters::SGD::Instance<SPEC_1>::USE_MASTER_PARAMETERS){
+            acc += abs_diff(device, p1.master_parameters, p2.parameters);
+        }
+        else if constexpr(nn::parameters::SGD::Instance<SPEC_2>::USE_MASTER_PARAMETERS){
+            acc += abs_diff(device, p1.parameters, p2.master_parameters);
+        }
+        return acc;
+    }
+    template<typename DEVICE, typename SPEC, typename MODE = Mode<mode::Default<>>>
+    RL_TOOLS_FUNCTION_PLACEMENT bool is_nan(DEVICE& device, const nn::parameters::SGD::Instance<SPEC>& p, const Mode<MODE>& mode = {}){
+        bool upstream_nan = is_nan(device, static_cast<const nn::parameters::Gradient::Instance<SPEC>&>(p), mode);
+        if constexpr(mode::is<MODE, nn::parameters::mode::ParametersOnly>){
+            return upstream_nan;
+        }
+        bool downstream_nan = is_nan(device, p.velocity, mode);
+        if constexpr(nn::parameters::SGD::Instance<SPEC>::USE_MASTER_PARAMETERS){
+            downstream_nan = downstream_nan || is_nan(device, p.master_parameters, mode);
+        }
+        return upstream_nan || downstream_nan;
+    }
+    template<typename DEVICE, typename SPEC, typename OPTIMIZER_SPEC, typename SELECTOR>
+    RL_TOOLS_FUNCTION_PLACEMENT void update(DEVICE& device, nn::parameters::SGD::Instance<SPEC>& parameter, nn::optimizers::Selected<nn::optimizers::SGD<OPTIMIZER_SPEC>, SELECTOR>& selection){
+        if constexpr(SELECTOR::template MATCH<nn::optimizers::UpdateTarget<SPEC>>){
+            update(device, parameter, selection.optimizer);
+        }
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/nn/optimizers/sgd/instance/persist.h

@@ -0,0 +1,37 @@
+#include "../../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_OPTIMIZERS_SGD_INSTANCE_PERSIST_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_OPTIMIZERS_SGD_INSTANCE_PERSIST_H
+
+#include "../sgd.h"
+#include "../../../parameters/persist.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template<typename DEVICE, typename CONTAINER, typename GROUP>
+    void save(DEVICE& device, nn::parameters::SGD::Instance<CONTAINER>& parameter, GROUP& group) {
+        save(device, (nn::parameters::Gradient::Instance<CONTAINER>&)parameter, group);
+        save(device, parameter.velocity, group, "velocity");
+        if constexpr(nn::parameters::SGD::Instance<CONTAINER>::USE_MASTER_PARAMETERS){
+            save(device, parameter.master_parameters, group, "master_parameters");
+        }
+    }
+    template<typename DEVICE, typename CONTAINER, typename GROUP>
+    bool load(DEVICE& device, nn::parameters::SGD::Instance<CONTAINER>& parameter, GROUP& group) {
+        bool success = load(device, (nn::parameters::Gradient::Instance<CONTAINER>&)parameter, group);
+        success &= load(device, parameter.velocity, group, "velocity");
+        if constexpr(nn::parameters::SGD::Instance<CONTAINER>::USE_MASTER_PARAMETERS){
+            if(group_exists(device, group, "master_parameters")){
+                success &= load(device, parameter.master_parameters, group, "master_parameters");
+            }
+            else{
+                // Backward compatibility with checkpoints generated before master parameters existed.
+                copy(device, device, parameter.parameters, parameter.master_parameters);
+            }
+        }
+        return success;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/nn/optimizers/sgd/instance/persist_code.h

@@ -0,0 +1,72 @@
+#include "../../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_OPTIMIZERS_SGD_INSTANCE_PERSIST_CODE_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_OPTIMIZERS_SGD_INSTANCE_PERSIST_CODE_H
+
+#include "../sgd.h"
+#include "../../../../nn/parameters/persist_code.h"
+
+#include <string>
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    std::string get_type_string(nn::parameters::SGD p){
+        return "RL_TOOLS""_NAMESPACE_WRAPPER ::rl_tools::nn::parameters::SGD";
+    }
+    template<typename DEVICE, typename SPEC>
+    persist::Code save_code_split(DEVICE& device, nn::parameters::SGD::Instance<SPEC>& parameter, std::string name, bool const_declaration=true, typename DEVICE::index_t indent=0, bool output_memory_only=false){
+        using TI = typename DEVICE::index_t;
+        std::stringstream indent_ss;
+        for(TI i=0; i < indent; i++){
+            indent_ss << "    ";
+        }
+        std::string ind = indent_ss.str();
+        std::stringstream ss, ss_header;
+        ss_header << nn::group::CodeGroups<typename nn::group::MembershipOf<SPEC>::TYPE>::headers();
+        auto plain = save_code_split(device, (nn::parameters::Gradient::Instance<SPEC>&) parameter, name, const_declaration, indent, true);
+        ss_header << plain.header;
+        ss_header << "#include <rl_tools/nn/optimizers/sgd/sgd.h>\n";
+        ss << plain.body;
+        ss << ind << "namespace " << name << " {\n";
+        auto velocity = save_code_split(device, parameter.velocity, "velocity_memory", const_declaration, indent+1);
+        ss_header << velocity.header;
+        ss << velocity.body;
+        if constexpr(nn::parameters::SGD::Instance<SPEC>::USE_MASTER_PARAMETERS){
+            auto master_parameters = save_code_split(device, parameter.master_parameters, "master_parameters_memory", const_declaration, indent+1);
+            ss_header << master_parameters.header;
+            ss << master_parameters.body;
+        }
+        if(!output_memory_only){
+            ss << ind << "    " << "using TYPE_POLICY = " << to_string(typename SPEC::TYPE_POLICY{}) << ";\n";
+            ss << ind << "    " << "using PARAMETER_SPEC = " << nn::group::CodeGroups<typename nn::group::MembershipOf<SPEC>::TYPE>::parameter_type("RL_TOOLS""_NAMESPACE_WRAPPER ::rl_tools::nn::parameters::SGD") << "::Specification<TYPE_POLICY, typename parameters_memory::SPEC::TI, typename parameters_memory::SPEC::SHAPE, "
+               << get_type_string_tag(device, typename SPEC::GROUP_TAG{})
+               << ", "
+               << get_type_string_tag(device, typename SPEC::CATEGORY_TAG{})
+               << ", true, true>;\n";
+            if constexpr(nn::parameters::SGD::Instance<SPEC>::USE_MASTER_PARAMETERS){
+                ss << ind << "    " << (const_declaration ? "constexpr " : "") << "RL_TOOLS""_NAMESPACE_WRAPPER ::rl_tools::nn::parameters::SGD::Instance<PARAMETER_SPEC> parameters = {{{parameters_memory::container}, gradient_memory::container}, velocity_memory::container, master_parameters_memory::container};\n";
+            }
+            else{
+                ss << ind << "    " << (const_declaration ? "constexpr " : "") << "RL_TOOLS""_NAMESPACE_WRAPPER ::rl_tools::nn::parameters::SGD::Instance<PARAMETER_SPEC> parameters = {{{parameters_memory::container}, gradient_memory::container}, velocity_memory::container};\n";
+            }
+        }
+        ss << ind << "}\n";
+        return {ss_header.str(), ss.str()};
+    }
+    template <typename DEVICE, typename SPEC>
+    std::string nn_analytics(DEVICE& device, nn::parameters::SGD::Instance<SPEC>& p) {
+        std::string data;
+        data += "{";
+        data += nn_analytics(device, static_cast<nn::parameters::Gradient::Instance<SPEC>&>(p), true) + ", ";
+        data += "\"velocity\": " + json(device, p.velocity);
+        if constexpr(nn::parameters::SGD::Instance<SPEC>::USE_MASTER_PARAMETERS){
+            data += ", \"master_parameters\": " + json(device, p.master_parameters);
+        }
+        data += "}";
+        return data;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/nn/optimizers/sgd/operations_cuda.h

@@ -0,0 +1,49 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_OPTIMIZERS_SGD_OPERATIONS_CUDA_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_OPTIMIZERS_SGD_OPERATIONS_CUDA_H
+#include "sgd.h"
+#include "../../../nn/parameters/operations_generic.h"
+#include "operations_generic.h"
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    namespace nn::optimizers::sgd::kernels{
+        template<typename DEVICE, typename SPEC>
+        __global__
+        void init(DEVICE device, nn::optimizers::SGD<SPEC> optimizer) {
+            static_assert(DEVICE::TAG && DEVICE::KERNEL);
+            typename nn::optimizers::SGD<SPEC>::PARAMETERS parameters = {
+                SPEC::DEFAULT_PARAMETERS::LEARNING_RATE,
+                SPEC::DEFAULT_PARAMETERS::MOMENTUM,
+                SPEC::DEFAULT_PARAMETERS::WEIGHT_DECAY,
+                SPEC::DEFAULT_PARAMETERS::NESTEROV
+            };
+            set(device, optimizer.parameters, parameters, 0);
+        }
+    }
+    template<typename DEV_SPEC, typename SPEC>
+    void init(devices::CUDA<DEV_SPEC>& device, nn::optimizers::SGD<SPEC>& optimizer) {
+        dim3 grid(1);
+        dim3 block(1);
+        devices::cuda::TAG<devices::CUDA<DEV_SPEC>, true> tag_device{};
+        nn::optimizers::sgd::kernels::init<<<grid, block, 0, device.stream>>>(tag_device, optimizer);
+        check_status(device);
+    }
+    template<typename DEV_SPEC, typename SPEC, typename MODEL>
+    void reset_optimizer_state(devices::CUDA<DEV_SPEC>& device, nn::optimizers::SGD<SPEC>& optimizer, MODEL& model) {
+        _reset_optimizer_state(device, model, optimizer);
+    }
+    template<typename DEV_SPEC, typename SPEC, typename MODEL, typename MODE = mode::Default<>>
+    void step(devices::CUDA<DEV_SPEC>& device, nn::optimizers::SGD<SPEC>& optimizer, MODEL& model, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        using SELECTOR = typename nn::optimizers::UpdateSelection<MODE>::TYPE;
+        if constexpr(utils::typing::is_same_v<SELECTOR, nn::group::All>){
+            update(device, model, optimizer);
+        }
+        else{
+            nn::optimizers::Selected<nn::optimizers::SGD<SPEC>, SELECTOR> selection{optimizer};
+            update(device, model, selection);
+        }
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/nn/optimizers/sgd/operations_generic.h

@@ -0,0 +1,50 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_OPTIMIZERS_SGD_OPERATIONS_GENERIC_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_OPTIMIZERS_SGD_OPERATIONS_GENERIC_H
+
+#include "sgd.h"
+#include "../../../nn/parameters/operations_generic.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::optimizers::SGD<SPEC>& optimizer){
+        malloc(device, optimizer.parameters);
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::optimizers::SGD<SPEC>& optimizer){
+        free(device, optimizer.parameters);
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void init(DEVICE& device, nn::optimizers::SGD<SPEC>& optimizer){
+        typename nn::optimizers::SGD<SPEC>::PARAMETERS parameters = {
+            SPEC::DEFAULT_PARAMETERS::LEARNING_RATE,
+            SPEC::DEFAULT_PARAMETERS::MOMENTUM,
+            SPEC::DEFAULT_PARAMETERS::WEIGHT_DECAY,
+            SPEC::DEFAULT_PARAMETERS::NESTEROV
+        };
+        set(device, optimizer.parameters, parameters, 0);
+    }
+    template<typename DEVICE, typename SPEC, typename MODEL>
+    RL_TOOLS_FUNCTION_PLACEMENT void reset_optimizer_state(DEVICE& device, nn::optimizers::SGD<SPEC>& optimizer, MODEL& model) {
+        _reset_optimizer_state(device, model, optimizer);
+    }
+    template<typename DEVICE, typename SPEC, typename MODEL, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void step(DEVICE& device, nn::optimizers::SGD<SPEC>& optimizer, MODEL& model, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        using SELECTOR = typename nn::optimizers::UpdateSelection<MODE>::TYPE;
+        if constexpr(utils::typing::is_same_v<SELECTOR, nn::group::All>){
+            update(device, model, optimizer);
+        }
+        else{
+            nn::optimizers::Selected<nn::optimizers::SGD<SPEC>, SELECTOR> selection{optimizer};
+            update(device, model, selection);
+        }
+    }
+    template<typename SOURCE_DEVICE, typename TARGET_DEVICE, typename SOURCE_SPEC, typename TARGET_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void copy(SOURCE_DEVICE& source_device, TARGET_DEVICE& target_device, const nn::optimizers::SGD<SOURCE_SPEC>& source, nn::optimizers::SGD<TARGET_SPEC>& target){
+        copy(source_device, target_device, source.parameters, target.parameters);
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/nn/optimizers/sgd/sgd.h

@@ -0,0 +1,92 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_OPTIMIZERS_SGD_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_OPTIMIZERS_SGD_H
+
+#include "../../../nn/parameters/parameters.h"
+#include "../update.h"
+#include "../../../utils/generic/typing.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::nn::optimizers{
+    namespace sgd{
+        template<typename T_T>
+        struct Parameters{
+            using T = T_T;
+            T learning_rate;
+            T momentum;
+            T weight_decay;
+            bool nesterov;
+        };
+        template <typename T_TYPE_POLICY>
+        struct DefaultParameters{
+            using T = typename T_TYPE_POLICY::DEFAULT;
+            static constexpr T LEARNING_RATE = 0.1;
+            static constexpr T MOMENTUM = 0.9;
+            static constexpr T WEIGHT_DECAY = 0;
+            static constexpr bool NESTEROV = false;
+            static constexpr bool ENABLE_WEIGHT_DECAY = false;
+        };
+        template <typename T_TYPE_POLICY, typename T_TI, typename T_DEFAULT_PARAMETERS = DefaultParameters<T_TYPE_POLICY>, bool T_DYNAMIC_ALLOCATION = true>
+        struct Specification{
+            using T = typename T_TYPE_POLICY::DEFAULT;
+            using TI = T_TI;
+            using DEFAULT_PARAMETERS = T_DEFAULT_PARAMETERS;
+            static constexpr bool ENABLE_WEIGHT_DECAY = DEFAULT_PARAMETERS::ENABLE_WEIGHT_DECAY;
+            static constexpr bool DYNAMIC_ALLOCATION = T_DYNAMIC_ALLOCATION;
+        };
+    }
+    template<typename T_SPEC>
+    struct SGD{
+        using SPEC = T_SPEC;
+        using T = typename SPEC::T;
+        using TI = typename SPEC::TI;
+        using DEFAULT_PARAMETERS = typename SPEC::DEFAULT_PARAMETERS;
+        using PARAMETERS = sgd::Parameters<T>;
+        Tensor<tensor::Specification<PARAMETERS, TI, tensor::Shape<TI, 1>, SPEC::DYNAMIC_ALLOCATION>> parameters;
+    };
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::nn::parameters{
+    struct SGD{
+        template <typename T_TYPE_POLICY, typename T_TI, typename T_SHAPE, typename T_GROUP_TAG, typename T_CATEGORY_TAG, bool T_DYNAMIC_ALLOCATION, bool T_CONST=false, typename T_NUMERIC_CATEGORY = numeric_types::categories::Parameter>
+        struct Specification{
+            using TYPE_POLICY = T_TYPE_POLICY;
+            using TI = T_TI;
+            using SHAPE = T_SHAPE;
+            using GROUP_TAG = T_GROUP_TAG;
+            using CATEGORY_TAG = T_CATEGORY_TAG;
+            static constexpr bool DYNAMIC_ALLOCATION = T_DYNAMIC_ALLOCATION;
+            static constexpr bool CONST = T_CONST;
+            using NUMERIC_CATEGORY = T_NUMERIC_CATEGORY;
+        };
+        template <typename T_SPEC, bool T_USE_MASTER_PARAMETERS = T_SPEC::TYPE_POLICY::template IS_SET<numeric_types::categories::MasterParameter> && !utils::typing::is_same_v<typename Gradient::Instance<T_SPEC>::PARENT::T_PARAMETER, typename T_SPEC::TYPE_POLICY::template GET<numeric_types::categories::MasterParameter>>>
+        struct Instance;
+        template <typename T_SPEC>
+        struct Instance<T_SPEC, false>: Gradient::Instance<T_SPEC>{
+            using SPEC = T_SPEC;
+            using T_PARAMETER = typename Gradient::Instance<T_SPEC>::PARENT::T_PARAMETER;
+            using T_MASTER_PARAMETER = typename T_SPEC::TYPE_POLICY::template GET<numeric_types::categories::MasterParameter>;
+            static constexpr bool USE_MASTER_PARAMETERS = false;
+            using T_VELOCITY = typename T_SPEC::TYPE_POLICY::template GET<numeric_types::categories::OptimizerState>;
+            using TENSOR_SPEC = tensor::Specification<T_VELOCITY, typename SPEC::TI, typename SPEC::SHAPE, SPEC::DYNAMIC_ALLOCATION, tensor::RowMajorStride<typename SPEC::SHAPE>, SPEC::CONST>;
+            Tensor<TENSOR_SPEC> velocity;
+        };
+        template <typename T_SPEC>
+        struct Instance<T_SPEC, true>: Gradient::Instance<T_SPEC>{
+            using SPEC = T_SPEC;
+            using T_PARAMETER = typename Gradient::Instance<T_SPEC>::PARENT::T_PARAMETER;
+            using T_MASTER_PARAMETER = typename T_SPEC::TYPE_POLICY::template GET<numeric_types::categories::MasterParameter>;
+            static constexpr bool USE_MASTER_PARAMETERS = true;
+            using T_VELOCITY = typename T_SPEC::TYPE_POLICY::template GET<numeric_types::categories::OptimizerState>;
+            using TENSOR_SPEC = tensor::Specification<T_VELOCITY, typename SPEC::TI, typename SPEC::SHAPE, SPEC::DYNAMIC_ALLOCATION, tensor::RowMajorStride<typename SPEC::SHAPE>, SPEC::CONST>;
+            Tensor<TENSOR_SPEC> velocity;
+            using MASTER_TENSOR_SPEC = tensor::Specification<T_MASTER_PARAMETER, typename SPEC::TI, typename SPEC::SHAPE, SPEC::DYNAMIC_ALLOCATION, tensor::RowMajorStride<typename SPEC::SHAPE>, SPEC::CONST>;
+            Tensor<MASTER_TENSOR_SPEC> master_parameters;
+        };
+    };
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/nn/optimizers/update.h

@@ -0,0 +1,54 @@
+#include "../../version.h"
+#include "../../rl_tools.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_OPTIMIZERS_UPDATE_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_OPTIMIZERS_UPDATE_H
+
+#include "../group.h"
+#include "../../mode/mode.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::nn::optimizers{
+    template<typename T_PARAMETER_SPEC>
+    struct UpdateTarget{
+        using PARAMETER_SPEC = T_PARAMETER_SPEC;
+        using MODEL_GROUPS = typename group::MembershipOf<PARAMETER_SPEC>::TYPE;
+        using GROUP_TAG = typename PARAMETER_SPEC::GROUP_TAG;
+        using CATEGORY_TAG = typename PARAMETER_SPEC::CATEGORY_TAG;
+    };
+
+    namespace mode{
+        template<typename T_BASE = rl_tools::mode::Default<>, typename T_SPEC = group::All>
+        struct Update: T_BASE{
+            using BASE = T_BASE;
+            using SPEC = T_SPEC;
+        };
+        template<typename T_BASE = rl_tools::mode::Default<>, typename T_SPEC = group::All>
+        struct Freeze: T_BASE{
+            using BASE = T_BASE;
+            using SPEC = T_SPEC;
+        };
+    }
+
+    template<typename MODE>
+    struct UpdateSelection: UpdateSelection<typename MODE::BASE>{};
+    template<>
+    struct UpdateSelection<rl_tools::mode::Final>{
+        using TYPE = group::All;
+    };
+    template<typename BASE, typename SELECTOR>
+    struct UpdateSelection<mode::Update<BASE, SELECTOR>>{
+        using TYPE = SELECTOR;
+    };
+    template<typename BASE, typename SELECTOR>
+    struct UpdateSelection<mode::Freeze<BASE, SELECTOR>>{
+        using TYPE = group::Not<SELECTOR>;
+    };
+
+    template<typename OPTIMIZER, typename SELECTOR>
+    struct Selected{
+        OPTIMIZER& optimizer;
+    };
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/nn/parameters/operations_generic.h

@@ -0,0 +1,98 @@
+#include "../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_PARAMETERS_OPERATIONS_GENERIC_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_PARAMETERS_OPERATIONS_GENERIC_H
+
+#include "parameters.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template <typename DEVICE, typename CONTAINER>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::parameters::Plain::Instance<CONTAINER>& p){
+        malloc(device, p.parameters);
+    }
+    template <typename DEVICE, typename CONTAINER>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::parameters::Plain::Instance<CONTAINER>& p){
+        free(device, p.parameters);
+    }
+    template <typename DEVICE, typename CONTAINER>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::parameters::Gradient::Instance<CONTAINER>& p){
+        malloc(device, (nn::parameters::Plain::Instance<CONTAINER>&) p);
+        malloc(device, p.gradient);
+    }
+    template <typename DEVICE, typename CONTAINER>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::parameters::Gradient::Instance<CONTAINER>& p){
+        free(device, (nn::parameters::Plain::Instance<CONTAINER>&) p);
+        free(device, p.gradient);
+    }
+    template<typename DEVICE, typename CONTAINER>
+    RL_TOOLS_FUNCTION_PLACEMENT void zero_gradient(DEVICE& device, nn::parameters::Gradient::Instance<CONTAINER>& container) {
+        set_all(device, container.gradient, 0);
+    }
+    template<typename DEVICE, typename SOURCE_SPEC, typename TARGET_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void add_gradient(DEVICE& device, nn::parameters::Gradient::Instance<SOURCE_SPEC>& source, nn::parameters::Gradient::Instance<TARGET_SPEC>& target) {
+        add(device, source.gradient, target.gradient);
+    }
+    template<typename SOURCE_DEVICE, typename TARGET_DEVICE, typename SOURCE_SPEC, typename TARGET_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void copy_gradient(SOURCE_DEVICE& source_device, TARGET_DEVICE& target_device, const nn::parameters::Gradient::Instance<SOURCE_SPEC>& source, nn::parameters::Gradient::Instance<TARGET_SPEC>& target) {
+        copy(source_device, target_device, source.gradient, target.gradient);
+    }
+
+    template<typename SOURCE_DEVICE, typename TARGET_DEVICE, typename SOURCE_SPEC, typename TARGET_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void copy(SOURCE_DEVICE& source_device, TARGET_DEVICE& target_device, const nn::parameters::Plain::Instance<SOURCE_SPEC>& source, nn::parameters::Plain::Instance<TARGET_SPEC>& target){
+        copy(source_device, target_device, source.parameters, target.parameters);
+    }
+
+    template<typename SOURCE_DEVICE, typename TARGET_DEVICE, typename SOURCE_SPEC, typename TARGET_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void copy(SOURCE_DEVICE& source_device, TARGET_DEVICE& target_device, const nn::parameters::Gradient::Instance<SOURCE_SPEC>& source, nn::parameters::Gradient::Instance<TARGET_SPEC>& target){
+        copy(source_device, target_device, (nn::parameters::Plain::Instance<SOURCE_SPEC>&) source, (nn::parameters::Plain::Instance<TARGET_SPEC>&) target);
+        copy(source_device, target_device, source.gradient, target.gradient);
+    }
+
+    template<typename SOURCE_DEVICE, typename TARGET_DEVICE, typename SOURCE, typename TARGET_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void copy_from_generic(SOURCE_DEVICE& source_device, TARGET_DEVICE& target_device, const SOURCE& source, nn::parameters::Plain::Instance<TARGET_SPEC>& target){
+        copy_from_generic(source_device, target_device, source.parameters, target.parameters);
+    }
+
+    template<typename SOURCE_DEVICE, typename TARGET_DEVICE, typename SOURCE, typename TARGET_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void copy_from_generic(SOURCE_DEVICE& source_device, TARGET_DEVICE& target_device, const SOURCE& source, nn::parameters::Gradient::Instance<TARGET_SPEC>& target){
+        copy_from_generic(source_device, target_device, static_cast<typename SOURCE::PARENT&>(source), static_cast<nn::parameters::Plain::Instance<TARGET_SPEC>&>(target));
+        copy_from_generic(source_device, target_device, source.gradient, target.gradient);
+    }
+
+    template<typename DEVICE, typename SPEC_1, typename SPEC_2>
+    RL_TOOLS_FUNCTION_PLACEMENT auto abs_diff(DEVICE& device, const nn::parameters::Plain::Instance<SPEC_1>& p1, const nn::parameters::Plain::Instance<SPEC_2>& p2){
+        typename decltype(p1.parameters)::SPEC::T acc = 0;
+        acc += abs_diff(device, p1.parameters, p2.parameters);
+        return acc;
+    }
+
+    template<typename DEVICE, typename SPEC_1, typename SPEC_2>
+    RL_TOOLS_FUNCTION_PLACEMENT auto abs_diff(DEVICE& device, const nn::parameters::Gradient::Instance<SPEC_1>& p1, const nn::parameters::Gradient::Instance<SPEC_2>& p2){
+        typename decltype(p1.parameters)::SPEC::T acc = 0;
+        acc += abs_diff(device, (nn::parameters::Plain::Instance<SPEC_1>&) p1, (nn::parameters::Plain::Instance<SPEC_2>&) p2);
+        acc += abs_diff(device, p1.gradient, p2.gradient);
+        return acc;
+    }
+
+    template<typename DEVICE, typename SPEC, typename MODE = Mode<mode::Default<>>>
+    RL_TOOLS_FUNCTION_PLACEMENT bool is_nan(DEVICE& device, const nn::parameters::Plain::Instance<SPEC>& p, const Mode<MODE>& mode = {}){
+        return is_nan(device, p.parameters, mode);
+    }
+
+
+    template<typename DEVICE, typename SPEC, typename MODE = Mode<mode::Default<>>>
+    RL_TOOLS_FUNCTION_PLACEMENT bool is_nan(DEVICE& device, const nn::parameters::Gradient::Instance<SPEC>& p, const Mode<MODE>& mode = {}){
+        bool upstream_nan = is_nan(device, static_cast<const nn::parameters::Plain::Instance<SPEC>&>(p), mode);
+        if constexpr(mode::is<MODE, nn::parameters::mode::ParametersOnly>){
+            return upstream_nan;
+        }
+        return upstream_nan || is_nan(device, p.gradient, mode);
+    }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT auto gradient_norm(DEVICE& device, const nn::parameters::Gradient::Instance<SPEC>& parameter){
+        return squared_sum(device, parameter.gradient);
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif