hyperray 0.0.1__py3-none-any.whl

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

@@ -0,0 +1,80 @@
+#include "../../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_OPTIMIZERS_ADAM_INSTANCE_PERSIST_CODE_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_OPTIMIZERS_ADAM_INSTANCE_PERSIST_CODE_H
+
+#include "../adam.h"
+#include "../../../../nn/parameters/persist_code.h"
+
+#include <string>
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    std::string get_type_string(nn::parameters::Adam p){
+        return "RL_TOOLS""_NAMESPACE_WRAPPER ::rl_tools::nn::parameters::Adam";
+    }
+    template<typename DEVICE, typename SPEC>
+    persist::Code save_code_split(DEVICE& device, nn::parameters::Adam::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/adam/adam.h>\n";
+        ss << plain.body;
+        ss << ind << "namespace " << name << " {\n";
+        auto gradient_first_order_moment = save_code_split(device, parameter.gradient_first_order_moment, "gradient_first_order_moment_memory", const_declaration, indent+1);
+        ss_header << gradient_first_order_moment.header;
+        ss << gradient_first_order_moment.body;
+        auto gradient_second_order_moment = save_code_split(device, parameter.gradient_second_order_moment, "gradient_second_order_moment_memory", const_declaration, indent+1);
+        ss_header << gradient_second_order_moment.header;
+        ss << gradient_second_order_moment.body;
+        if constexpr(nn::parameters::Adam::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 << "    " << "static_assert(RL_TOOLS""_NAMESPACE_WRAPPER ::rl_tools::utils::typing::is_same_v<parameters_memory::CONTAINER_TYPE, gradient_memory::CONTAINER_TYPE>);\n";
+            // ss << ind << "    " << "static_assert(RL_TOOLS""_NAMESPACE_WRAPPER ::rl_tools::utils::typing::is_same_v<gradient_memory::CONTAINER_TYPE, gradient_first_order_moment_memory::CONTAINER_TYPE>);\n";
+            // ss << ind << "    " << "static_assert(RL_TOOLS""_NAMESPACE_WRAPPER ::rl_tools::utils::typing::is_same_v<gradient_memory::CONTAINER_TYPE, gradient_second_order_moment_memory::CONTAINER_TYPE>);\n";
+            // ss << ind << "    " << "using PARAMETER_SPEC = " << "RL_TOOLS""_NAMESPACE_WRAPPER ::rl_tools::nn::parameters::Adam::Specification<parameters_memory::CONTAINER_TYPE, "
+            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::Adam") << "::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::Adam::Instance<SPEC>::USE_MASTER_PARAMETERS){
+                ss << ind << "    " << (const_declaration ? "constexpr " : "") << "RL_TOOLS""_NAMESPACE_WRAPPER ::rl_tools::nn::parameters::Adam::Instance<PARAMETER_SPEC> parameters = {{{parameters_memory::container}, gradient_memory::container}, gradient_first_order_moment_memory::container, gradient_second_order_moment_memory::container, master_parameters_memory::container};\n";
+            }
+            else{
+                ss << ind << "    " << (const_declaration ? "constexpr " : "") << "RL_TOOLS""_NAMESPACE_WRAPPER ::rl_tools::nn::parameters::Adam::Instance<PARAMETER_SPEC> parameters = {{{parameters_memory::container}, gradient_memory::container}, gradient_first_order_moment_memory::container, gradient_second_order_moment_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::Adam::Instance<SPEC>& p) {
+        std::string data;
+        data += "{";
+        data += nn_analytics(device, static_cast<nn::parameters::Gradient::Instance<SPEC>&>(p), true) + ", ";
+        data += "\"gradient_first_order_moment\": " + json(device, p.gradient_first_order_moment) + ", ";
+        data += "\"gradient_second_order_moment\": " + json(device, p.gradient_second_order_moment);
+        if constexpr(nn::parameters::Adam::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/adam/operations_cuda.h

@@ -0,0 +1,85 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_OPTIMIZERS_ADAM_OPERATIONS_CUDA_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_OPTIMIZERS_ADAM_OPERATIONS_CUDA_H
+
+#include "adam.h"
+#include "../../../nn/layers/dense/layer.h"
+#include "../../../nn/parameters/operations_generic.h"
+#include "../../../utils/polyak/operations_generic.h"
+
+#include "operations_generic.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    namespace nn::optimizers::adam::kernels{
+        template<typename DEVICE, typename SPEC>
+        __global__
+        void init(DEVICE device, nn::optimizers::Adam<SPEC> optimizer) {
+            static_assert(DEVICE::TAG && DEVICE::KERNEL);
+            typename nn::optimizers::Adam<SPEC>::PARAMETERS parameters = {
+                SPEC::DEFAULT_PARAMETERS::ALPHA,
+                SPEC::DEFAULT_PARAMETERS::BETA_1,
+                SPEC::DEFAULT_PARAMETERS::BETA_2,
+                SPEC::DEFAULT_PARAMETERS::EPSILON,
+                SPEC::DEFAULT_PARAMETERS::EPSILON_SQRT,
+                SPEC::DEFAULT_PARAMETERS::WEIGHT_DECAY,
+                SPEC::DEFAULT_PARAMETERS::WEIGHT_DECAY_INPUT,
+                SPEC::DEFAULT_PARAMETERS::WEIGHT_DECAY_OUTPUT,
+                SPEC::DEFAULT_PARAMETERS::BIAS_LR_FACTOR
+            };
+            set(device, optimizer.parameters, parameters, 0);
+        }
+        template<typename DEVICE, typename SPEC>
+        __global__
+        void reset_optimizer_state(DEVICE device, nn::optimizers::Adam<SPEC> optimizer) {
+            static_assert(DEVICE::TAG && DEVICE::KERNEL);
+            set(device, optimizer.age, 1, 0);
+        }
+    }
+    template<typename DEV_SPEC, typename SPEC>
+    void init(devices::CUDA<DEV_SPEC>& device, nn::optimizers::Adam<SPEC>& optimizer) {
+        dim3 activation_grid(1);
+        dim3 activation_block(1);
+        devices::cuda::TAG<devices::CUDA<DEV_SPEC>, true> tag_device{};
+        nn::optimizers::adam::kernels::init<<<activation_grid, activation_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::Adam<SPEC>& optimizer, MODEL& model) {
+        dim3 activation_grid(1);
+        dim3 activation_block(1);
+        devices::cuda::TAG<devices::CUDA<DEV_SPEC>, true> tag_device{};
+        nn::optimizers::adam::kernels::reset_optimizer_state<<<activation_grid, activation_block, 0, device.stream>>>(tag_device, optimizer);
+        check_status(device);
+        _reset_optimizer_state(device, model, optimizer);
+    }
+    namespace nn::optimizers::adam::kernels{
+        template<typename DEVICE, typename SPEC>
+        __global__
+        void step(DEVICE device, nn::optimizers::Adam<SPEC> optimizer) {
+            static_assert(DEVICE::TAG && DEVICE::KERNEL);
+            _step(device, optimizer);
+        }
+    }
+    template<typename DEV_SPEC, typename SPEC, typename MODEL, typename MODE = mode::Default<>>
+    void step(devices::CUDA<DEV_SPEC>& device, nn::optimizers::Adam<SPEC>& optimizer, MODEL& model, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        using DEVICE = devices::CUDA<DEV_SPEC>;
+        dim3 activation_grid(1);
+        dim3 activation_block(1);
+        devices::cuda::TAG<devices::CUDA<DEV_SPEC>, true> tag_device{};
+        nn::optimizers::adam::kernels::step<<<activation_grid, activation_block, 0, device.stream>>>(tag_device, optimizer);
+        check_status(device);
+        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::Adam<SPEC>, SELECTOR> selection{optimizer};
+            update(device, model, selection);
+        }
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif

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

@@ -0,0 +1,109 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_OPTIMIZERS_ADAM_OPERATIONS_GENERIC_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_OPTIMIZERS_ADAM_OPERATIONS_GENERIC_H
+
+#include "adam.h"
+#include "../../../nn/layers/dense/layer.h"
+#include "../../../nn/parameters/operations_generic.h"
+#include "../../../utils/polyak/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::Adam<SPEC>& optimizer){
+        malloc(device, optimizer.age);
+        malloc(device, optimizer.first_order_moment_bias_correction);
+        malloc(device, optimizer.second_order_moment_bias_correction);
+        malloc(device, optimizer.parameters);
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::optimizers::Adam<SPEC>& optimizer){
+        free(device, optimizer.age);
+        free(device, optimizer.first_order_moment_bias_correction);
+        free(device, optimizer.second_order_moment_bias_correction);
+        free(device, optimizer.parameters);
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void init(DEVICE& device, nn::optimizers::Adam<SPEC>& optimizer){
+        typename nn::optimizers::Adam<SPEC>::PARAMETERS parameters = {
+            SPEC::DEFAULT_PARAMETERS::ALPHA,
+            SPEC::DEFAULT_PARAMETERS::BETA_1,
+            SPEC::DEFAULT_PARAMETERS::BETA_2,
+            SPEC::DEFAULT_PARAMETERS::EPSILON,
+            SPEC::DEFAULT_PARAMETERS::EPSILON_SQRT,
+            SPEC::DEFAULT_PARAMETERS::WEIGHT_DECAY,
+            SPEC::DEFAULT_PARAMETERS::WEIGHT_DECAY_INPUT,
+            SPEC::DEFAULT_PARAMETERS::WEIGHT_DECAY_OUTPUT,
+            SPEC::DEFAULT_PARAMETERS::BIAS_LR_FACTOR
+        };
+        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::Adam<SPEC>& optimizer, MODEL& model) {
+        set(device, optimizer.age, 1, 0);
+        _reset_optimizer_state(device, model, optimizer);
+    }
+
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void _step(DEVICE& device, nn::optimizers::Adam<SPEC>& optimizer){
+        using T = typename SPEC::T;
+        using TI = typename DEVICE::index_t;
+        TI age = get(device, optimizer.age, 0);
+        const auto& optimizer_parameters = get(device, optimizer.parameters, 0);
+        T first_order_moment_bias_correction  = 1/(1 - math::pow(device.math, optimizer_parameters.beta_1, (T)age));
+        set(device, optimizer.first_order_moment_bias_correction, first_order_moment_bias_correction, 0);
+        T second_order_moment_bias_correction = 1/(1 - math::pow(device.math, optimizer_parameters.beta_2, (T)age));
+        set(device, optimizer.second_order_moment_bias_correction, second_order_moment_bias_correction, 0);
+        set(device, optimizer.age, age + 1, 0);
+    }
+    template<typename DEVICE, typename SPEC, typename MODEL, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void step(DEVICE& device, nn::optimizers::Adam<SPEC>& optimizer, MODEL& model, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        _step(device, optimizer);
+        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::Adam<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::Adam<SOURCE_SPEC>& source, nn::optimizers::Adam<TARGET_SPEC>& target){
+        copy(source_device, target_device, source.age, target.age);
+        copy(source_device, target_device, source.parameters, target.parameters);
+    }
+
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT bool is_nan(DEVICE& device, const nn::parameters::Adam::Instance<SPEC>& p){
+        bool param_nan = is_nan(device, (nn::parameters::Gradient::Instance<SPEC>&) p);
+        return param_nan || is_nan(device, p.gradient_first_order_moment) || is_nan(device, p.gradient_second_order_moment);
+    }
+    template <typename DEVICE, typename T>
+    RL_TOOLS_FUNCTION_PLACEMENT T abs_diff(DEVICE& device, const nn::optimizers::adam::Parameters<T>& p1, const nn::optimizers::adam::Parameters<T>& p2){
+        T acc = 0;
+        acc += math::abs(device.math, p1.alpha - p2.alpha);
+        acc += math::abs(device.math, p1.beta_1 - p2.beta_1);
+        acc += math::abs(device.math, p1.beta_2 - p2.beta_2);
+        acc += math::abs(device.math, p1.epsilon - p2.epsilon);
+        acc += math::abs(device.math, p1.epsilon_sqrt - p2.epsilon_sqrt);
+        acc += math::abs(device.math, p1.weight_decay - p2.weight_decay);
+        acc += math::abs(device.math, p1.weight_decay_input - p2.weight_decay_input);
+        acc += math::abs(device.math, p1.weight_decay_output - p2.weight_decay_output);
+        acc += math::abs(device.math, p1.bias_lr_factor - p2.bias_lr_factor);
+        return acc;
+    }
+    template <typename DEVICE, typename SPEC_1, typename SPEC_2>
+    RL_TOOLS_FUNCTION_PLACEMENT typename SPEC_1::T abs_diff(DEVICE& device, nn::optimizers::Adam<SPEC_1>& o1, nn::optimizers::Adam<SPEC_2>& o2){
+        using T = typename SPEC_1::T;
+        T acc = 0;
+        acc += abs_diff(device, o1.age, o2.age);
+        acc += abs_diff(device, o1.first_order_moment_bias_correction, o2.first_order_moment_bias_correction);
+        acc += abs_diff(device, o1.second_order_moment_bias_correction, o2.second_order_moment_bias_correction);
+        acc += abs_diff(device, get(device, o1.parameters, 0), get(device, o2.parameters, 0));
+        return acc;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

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

@@ -0,0 +1,25 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_OPTIMIZERS_ADAM_PERSIST_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_OPTIMIZERS_ADAM_PERSIST_H
+#include "adam.h"
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template<typename DEVICE, typename SPEC, typename GROUP>
+    void save(DEVICE& device, nn::optimizers::Adam<SPEC>& optimizer, GROUP& group) {
+        save_binary(device, data(optimizer.parameters), 1, group, "parameters");
+        save(device, optimizer.age, group, "age");
+        save(device, optimizer.first_order_moment_bias_correction, group, "first_order_moment_bias_correction");
+        save(device, optimizer.second_order_moment_bias_correction, group, "second_order_moment_bias_correction");
+    }
+    template<typename DEVICE, typename SPEC, typename GROUP>
+    bool load(DEVICE& device, nn::optimizers::Adam<SPEC>& optimizer, GROUP& group) {
+        bool success = load_binary(device, data(optimizer.parameters), 1, group, "parameters");
+        success &= load(device, optimizer.age, group, "age");
+        success &= load(device, optimizer.first_order_moment_bias_correction, group, "first_order_moment_bias_correction");
+        success &= load(device, optimizer.second_order_moment_bias_correction, group, "second_order_moment_bias_correction");
+        return success;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

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

@@ -0,0 +1,107 @@
+#include "../../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_OPTIMIZERS_LAMB_INSTANCE_OPERATIONS_GENERIC_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_OPTIMIZERS_LAMB_INSTANCE_OPERATIONS_GENERIC_H
+
+#include "../lamb.h"
+#include "../../adam/instance/operations_generic.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template<typename DEVICE, typename SPEC, typename PARAMETER_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void _reset_optimizer_state(DEVICE& device, nn::parameters::Adam::Instance<PARAMETER_SPEC>& parameter, nn::optimizers::Lamb<SPEC>& optimizer){
+        _reset_optimizer_state(device, parameter, static_cast<nn::optimizers::Adam<SPEC>&>(optimizer));
+    }
+
+    template<typename DEVICE, typename SPEC, typename PARAMETER_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void gradient_descent(DEVICE& device, nn::parameters::Adam::Instance<PARAMETER_SPEC>& parameter, nn::optimizers::Lamb<SPEC>& optimizer){
+        using TI = typename DEVICE::index_t;
+        using T_OPTIMIZER = typename PARAMETER_SPEC::TYPE_POLICY::template GET<numeric_types::categories::OptimizerState>;
+        using T_PARAMETER = typename decltype(parameter.parameters)::T;
+        const auto& optimizer_parameters = get(device, optimizer.parameters, 0);
+        const auto& lamb_params = get(device, optimizer.lamb_parameters, 0);
+        auto parameters = matrix_view(device, parameter.parameters);
+        auto gradient_first_order_moment = matrix_view(device, parameter.gradient_first_order_moment);
+        auto gradient_second_order_moment = matrix_view(device, parameter.gradient_second_order_moment);
+        constexpr TI ROWS = decltype(parameters)::ROWS;
+        constexpr TI COLS = decltype(parameters)::COLS;
+
+        T_OPTIMIZER weight_norm_sq = 0;
+        T_OPTIMIZER update_norm_sq = 0;
+        for(TI row_i = 0; row_i < ROWS; row_i++){
+            for(TI col_i = 0; col_i < COLS; col_i++){
+                T_OPTIMIZER w = get(parameters, row_i, col_i);
+                weight_norm_sq += w * w;
+
+                T_OPTIMIZER pre_sqrt_term = get(gradient_second_order_moment, row_i, col_i) * get(device, optimizer.second_order_moment_bias_correction, 0);
+                pre_sqrt_term = math::max(device.math, pre_sqrt_term, (T_OPTIMIZER)optimizer_parameters.epsilon_sqrt);
+                T_OPTIMIZER r = get(device, optimizer.first_order_moment_bias_correction, 0) * get(gradient_first_order_moment, row_i, col_i) / (math::sqrt(device.math, pre_sqrt_term) + optimizer_parameters.epsilon);
+                if constexpr(utils::typing::is_same_v<typename PARAMETER_SPEC::CATEGORY_TAG, nn::parameters::categories::Weights>){
+                    if constexpr(utils::typing::is_same_v<typename PARAMETER_SPEC::GROUP_TAG, nn::parameters::groups::Normal> && SPEC::ENABLE_WEIGHT_DECAY){
+                        r += w * optimizer_parameters.weight_decay / 2;
+                    }
+                    if constexpr(utils::typing::is_same_v<typename PARAMETER_SPEC::GROUP_TAG, nn::parameters::groups::Input> && SPEC::ENABLE_WEIGHT_DECAY){
+                        r += w * optimizer_parameters.weight_decay_input / 2;
+                    }
+                    if constexpr(utils::typing::is_same_v<typename PARAMETER_SPEC::GROUP_TAG, nn::parameters::groups::Output> && SPEC::ENABLE_WEIGHT_DECAY){
+                        r += w * optimizer_parameters.weight_decay_output / 2;
+                    }
+                }
+                update_norm_sq += r * r;
+            }
+        }
+
+        T_OPTIMIZER weight_norm = math::sqrt(device.math, weight_norm_sq);
+        T_OPTIMIZER update_norm = math::sqrt(device.math, update_norm_sq);
+
+        T_OPTIMIZER trust_ratio = 1;
+        if(weight_norm > 0 && update_norm > 0){
+            trust_ratio = weight_norm / update_norm;
+            trust_ratio = math::min(device.math, trust_ratio, lamb_params.upper_bound_trust_ratio);
+            trust_ratio = math::max(device.math, trust_ratio, lamb_params.lower_bound_trust_ratio);
+        }
+
+        for(TI row_i = 0; row_i < ROWS; row_i++){
+            for(TI col_i = 0; col_i < COLS; col_i++){
+                T_OPTIMIZER pre_sqrt_term = get(gradient_second_order_moment, row_i, col_i) * get(device, optimizer.second_order_moment_bias_correction, 0);
+                pre_sqrt_term = math::max(device.math, pre_sqrt_term, (T_OPTIMIZER)optimizer_parameters.epsilon_sqrt);
+                T_OPTIMIZER r = get(device, optimizer.first_order_moment_bias_correction, 0) * get(gradient_first_order_moment, row_i, col_i) / (math::sqrt(device.math, pre_sqrt_term) + optimizer_parameters.epsilon);
+                if constexpr(utils::typing::is_same_v<typename PARAMETER_SPEC::CATEGORY_TAG, nn::parameters::categories::Biases> && SPEC::ENABLE_BIAS_LR_FACTOR){
+                    r *= optimizer_parameters.bias_lr_factor;
+                }
+                if constexpr(utils::typing::is_same_v<typename PARAMETER_SPEC::CATEGORY_TAG, nn::parameters::categories::Weights>){
+                    if constexpr(utils::typing::is_same_v<typename PARAMETER_SPEC::GROUP_TAG, nn::parameters::groups::Normal> && SPEC::ENABLE_WEIGHT_DECAY){
+                        r += get(parameters, row_i, col_i) * optimizer_parameters.weight_decay / 2;
+                    }
+                    if constexpr(utils::typing::is_same_v<typename PARAMETER_SPEC::GROUP_TAG, nn::parameters::groups::Input> && SPEC::ENABLE_WEIGHT_DECAY){
+                        r += get(parameters, row_i, col_i) * optimizer_parameters.weight_decay_input / 2;
+                    }
+                    if constexpr(utils::typing::is_same_v<typename PARAMETER_SPEC::GROUP_TAG, nn::parameters::groups::Output> && SPEC::ENABLE_WEIGHT_DECAY){
+                        r += get(parameters, row_i, col_i) * optimizer_parameters.weight_decay_output / 2;
+                    }
+                }
+                T_OPTIMIZER value = get(parameters, row_i, col_i);
+                value -= optimizer_parameters.alpha * trust_ratio * r;
+                set(parameters, row_i, col_i, (T_PARAMETER)value);
+            }
+        }
+    }
+
+    template<typename DEVICE, typename SPEC, typename LAMB_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void update(DEVICE& device, nn::parameters::Adam::Instance<SPEC>& parameter, nn::optimizers::Lamb<LAMB_SPEC>& optimizer){
+        using PARAMETERS = typename LAMB_SPEC::DEFAULT_PARAMETERS;
+        const auto& optimizer_parameters = get(device, optimizer.parameters, 0);
+        utils::polyak::update(device, parameter.gradient, parameter.gradient_first_order_moment, optimizer_parameters.beta_1, PARAMETERS::ENABLE_GRADIENT_CLIPPING, PARAMETERS::GRADIENT_CLIP_VALUE);
+        utils::polyak::update_squared(device, parameter.gradient, parameter.gradient_second_order_moment, optimizer_parameters.beta_2, PARAMETERS::ENABLE_GRADIENT_CLIPPING, PARAMETERS::GRADIENT_CLIP_VALUE);
+        gradient_descent(device, parameter, optimizer);
+    }
+    template<typename DEVICE, typename SPEC, typename OPTIMIZER_SPEC, typename SELECTOR>
+    RL_TOOLS_FUNCTION_PLACEMENT void update(DEVICE& device, nn::parameters::Adam::Instance<SPEC>& parameter, nn::optimizers::Selected<nn::optimizers::Lamb<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/lamb/lamb.h

@@ -0,0 +1,40 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_OPTIMIZERS_LAMB_LAMB_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_OPTIMIZERS_LAMB_LAMB_H
+
+#include "../adam/adam.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::nn::optimizers{
+    namespace lamb{
+        template<typename T_T>
+        struct Parameters{
+            using T = T_T;
+            T upper_bound_trust_ratio;
+            T lower_bound_trust_ratio;
+        };
+        template <typename TYPE_POLICY>
+        struct DEFAULT_PARAMETERS{
+            using T = typename TYPE_POLICY::DEFAULT;
+            static constexpr T UPPER_BOUND_TRUST_RATIO = 10.0;
+            static constexpr T LOWER_BOUND_TRUST_RATIO = 0.0;
+        };
+        template <typename T_TYPE_POLICY, typename T_TI, typename T_ADAM_DEFAULT_PARAMETERS = adam::DEFAULT_PARAMETERS_TENSORFLOW<T_TYPE_POLICY>, typename T_LAMB_DEFAULT_PARAMETERS = DEFAULT_PARAMETERS<T_TYPE_POLICY>, bool T_DYNAMIC_ALLOCATION = true>
+        struct Specification: adam::Specification<T_TYPE_POLICY, T_TI, T_ADAM_DEFAULT_PARAMETERS, T_DYNAMIC_ALLOCATION>{
+            using LAMB_DEFAULT_PARAMETERS = T_LAMB_DEFAULT_PARAMETERS;
+        };
+    }
+    template<typename T_SPEC>
+    struct Lamb: Adam<T_SPEC>{
+        using SPEC = T_SPEC;
+        using T = typename SPEC::T;
+        using TI = typename SPEC::TI;
+        using LAMB_DEFAULT_PARAMETERS = typename SPEC::LAMB_DEFAULT_PARAMETERS;
+        using LAMB_PARAMETERS = lamb::Parameters<T>;
+        Tensor<tensor::Specification<LAMB_PARAMETERS, TI, tensor::Shape<TI, 1>, SPEC::DYNAMIC_ALLOCATION>> lamb_parameters;
+    };
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif

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

@@ -0,0 +1,68 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_OPTIMIZERS_LAMB_OPERATIONS_GENERIC_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_OPTIMIZERS_LAMB_OPERATIONS_GENERIC_H
+
+#include "lamb.h"
+#include "../adam/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::Lamb<SPEC>& optimizer){
+        malloc(device, static_cast<nn::optimizers::Adam<SPEC>&>(optimizer));
+        malloc(device, optimizer.lamb_parameters);
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::optimizers::Lamb<SPEC>& optimizer){
+        free(device, static_cast<nn::optimizers::Adam<SPEC>&>(optimizer));
+        free(device, optimizer.lamb_parameters);
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void init(DEVICE& device, nn::optimizers::Lamb<SPEC>& optimizer){
+        init(device, static_cast<nn::optimizers::Adam<SPEC>&>(optimizer));
+        typename nn::optimizers::Lamb<SPEC>::LAMB_PARAMETERS lamb_params = {
+            SPEC::LAMB_DEFAULT_PARAMETERS::UPPER_BOUND_TRUST_RATIO,
+            SPEC::LAMB_DEFAULT_PARAMETERS::LOWER_BOUND_TRUST_RATIO,
+        };
+        set(device, optimizer.lamb_parameters, lamb_params, 0);
+    }
+    template<typename DEVICE, typename SPEC, typename MODEL>
+    RL_TOOLS_FUNCTION_PLACEMENT void reset_optimizer_state(DEVICE& device, nn::optimizers::Lamb<SPEC>& optimizer, MODEL& model){
+        set(device, optimizer.age, 1, 0);
+        _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::Lamb<SPEC>& optimizer, MODEL& model, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        _step(device, static_cast<nn::optimizers::Adam<SPEC>&>(optimizer));
+        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::Lamb<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::Lamb<SOURCE_SPEC>& source, nn::optimizers::Lamb<TARGET_SPEC>& target){
+        copy(source_device, target_device, static_cast<const nn::optimizers::Adam<SOURCE_SPEC>&>(source), static_cast<nn::optimizers::Adam<TARGET_SPEC>&>(target));
+        copy(source_device, target_device, source.lamb_parameters, target.lamb_parameters);
+    }
+    template <typename DEVICE, typename T>
+    RL_TOOLS_FUNCTION_PLACEMENT T abs_diff(DEVICE& device, const nn::optimizers::lamb::Parameters<T>& p1, const nn::optimizers::lamb::Parameters<T>& p2){
+        T acc = 0;
+        acc += math::abs(device.math, p1.upper_bound_trust_ratio - p2.upper_bound_trust_ratio);
+        acc += math::abs(device.math, p1.lower_bound_trust_ratio - p2.lower_bound_trust_ratio);
+        return acc;
+    }
+    template <typename DEVICE, typename SPEC_1, typename SPEC_2>
+    RL_TOOLS_FUNCTION_PLACEMENT typename SPEC_1::T abs_diff(DEVICE& device, nn::optimizers::Lamb<SPEC_1>& o1, nn::optimizers::Lamb<SPEC_2>& o2){
+        using T = typename SPEC_1::T;
+        T acc = abs_diff(device, static_cast<nn::optimizers::Adam<SPEC_1>&>(o1), static_cast<nn::optimizers::Adam<SPEC_2>&>(o2));
+        acc += abs_diff(device, get(device, o1.lamb_parameters, 0), get(device, o2.lamb_parameters, 0));
+        return acc;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

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

@@ -0,0 +1,3 @@
+#include "adam/adam.h"
+#include "lamb/lamb.h"
+#include "sgd/sgd.h"

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

@@ -0,0 +1,85 @@
+#include "../../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_OPTIMIZERS_SGD_INSTANCE_OPERATIONS_CUDA_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_OPTIMIZERS_SGD_INSTANCE_OPERATIONS_CUDA_H
+
+#include "../sgd.h"
+#include "operations_generic.h"
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools {
+    namespace nn::optimizers::sgd::cuda {
+        template<typename DEV_SPEC, typename PARAMETER_SPEC, typename SPEC>
+        __global__
+        void update_kernel(devices::CUDA<DEV_SPEC> device, nn::parameters::SGD::Instance<PARAMETER_SPEC> parameter, nn::optimizers::SGD<SPEC> optimizer) {
+            static_assert(DEV_SPEC::TAG && DEV_SPEC::KERNEL);
+            using DEVICE = devices::CUDA<DEV_SPEC>;
+            using TI = typename DEVICE::index_t;
+
+            const auto& optimizer_parameters = get_ref(device, optimizer.parameters, 0);
+
+            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_SPEC::TYPE_POLICY::template GET<numeric_types::categories::OptimizerState>;
+            using T_PARAMETER = typename decltype(parameter.parameters)::T;
+            using T_MASTER_PARAMETER = typename PARAMETER_INSTANCE::T_MASTER_PARAMETER;
+            auto params = matrix_view(device, parameter.parameters);
+            auto grad = matrix_view(device, parameter.gradient);
+            auto vel = matrix_view(device, parameter.velocity);
+            constexpr TI ROWS = decltype(params)::ROWS;
+            constexpr TI COLS = decltype(params)::COLS;
+
+            TI col_i = blockIdx.x * blockDim.x + threadIdx.x;
+            TI row_i = blockIdx.y * blockDim.y + threadIdx.y;
+            if(col_i < COLS && row_i < ROWS){
+                T_VELOCITY lr = (T_VELOCITY)optimizer_parameters.learning_rate;
+                T_VELOCITY mom = (T_VELOCITY)optimizer_parameters.momentum;
+                T_VELOCITY wd = (T_VELOCITY)optimizer_parameters.weight_decay;
+                T_VELOCITY value;
+                if constexpr(USE_MASTER_PARAMETERS){
+                    auto master_params = matrix_view(device, parameter.master_parameters);
+                    value = (T_VELOCITY)get(master_params, row_i, col_i);
+                }
+                else{
+                    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 * wd;
+                    }
+                }
+                T_VELOCITY v = mom * get(vel, row_i, col_i) + g;
+                set(vel, row_i, col_i, v);
+                T_VELOCITY param_update = optimizer_parameters.nesterov ? mom * v + g : v;
+                value -= lr * param_update;
+                if constexpr(USE_MASTER_PARAMETERS){
+                    auto master_params = matrix_view(device, parameter.master_parameters);
+                    set(master_params, row_i, col_i, (T_MASTER_PARAMETER)value);
+                }
+                set(params, row_i, col_i, (T_PARAMETER)value);
+            }
+        }
+    }
+    template<typename DEV_SPEC, typename SPEC, typename PARAMETERS>
+    void update(devices::CUDA<DEV_SPEC>& device, nn::parameters::SGD::Instance<SPEC>& p, nn::optimizers::SGD<PARAMETERS>& optimizer) {
+        constexpr typename devices::CUDA<DEV_SPEC>::index_t BLOCKSIZE_ROWS = 32;
+        constexpr typename devices::CUDA<DEV_SPEC>::index_t BLOCKSIZE_COLS = 32;
+        using MATRIX_SPEC = typename decltype(matrix_view(device, p.parameters))::SPEC;
+        constexpr typename devices::CUDA<DEV_SPEC>::index_t N_BLOCKS_ROWS = RL_TOOLS_DEVICES_CUDA_CEIL(MATRIX_SPEC::ROWS, BLOCKSIZE_ROWS);
+        constexpr typename devices::CUDA<DEV_SPEC>::index_t N_BLOCKS_COLS = RL_TOOLS_DEVICES_CUDA_CEIL(MATRIX_SPEC::COLS, BLOCKSIZE_COLS);
+        dim3 grid(N_BLOCKS_COLS, N_BLOCKS_ROWS);
+        dim3 block(BLOCKSIZE_COLS, BLOCKSIZE_ROWS);
+        devices::cuda::TAG<devices::CUDA<DEV_SPEC>, true> tag_device{};
+        nn::optimizers::sgd::cuda::update_kernel<<<grid, block, 0, device.stream>>>(tag_device, p, optimizer);
+        check_status(device);
+    }
+    template<typename DEV_SPEC, typename SPEC, typename OPTIMIZER_SPEC, typename SELECTOR>
+    void update(devices::CUDA<DEV_SPEC>& 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