hyperray 0.0.1__py3-none-any.whl

hyperdrone/_vendor/rl-tools/include/rl_tools/nn/layers/td3_sampling/layer.h

@@ -0,0 +1,123 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_LAYERS_TD3_SAMPLING_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_LAYERS_TD3_SAMPLING_H
+#include "../../../nn/activation_functions.h"
+#include "../../../utils/generic/typing.h"
+#include "../../../containers/matrix/matrix.h"
+#include "../../../nn/parameters/parameters.h"
+#include "../dense/layer.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    // This layer performs the sampling and squashing for SAC
+    namespace nn::layers::td3_sampling{
+        namespace mode{
+            template <typename T_BASE, typename T_SPEC = bool>
+            struct ExternalNoise{
+                using SPEC = T_SPEC;
+                using BASE = T_BASE;
+                // this mode uses the noise from the Buffer for debugging / no-side-effect inference
+            };
+        }
+        template <typename TYPE_POLICY>
+        struct DefaultParameters{
+            static constexpr typename TYPE_POLICY::DEFAULT STD = 0.1;
+        };
+
+        struct State{};
+
+        template <typename T_TI, typename T_SPEC, bool T_DYNAMIC_ALLOCATION>
+        struct BufferSpecification {
+            using TI = T_TI;
+            using SPEC = T_SPEC;
+            static constexpr bool DYNAMIC_ALLOCATION = T_DYNAMIC_ALLOCATION;
+        };
+
+        template <typename BUFFER_SPEC>
+        struct Buffer{
+            using LAYER_SPEC = typename BUFFER_SPEC::SPEC;
+            using T = typename LAYER_SPEC::TYPE_POLICY::DEFAULT;
+            using NOISE_CONTAINER_SPEC = matrix::Specification<T, typename BUFFER_SPEC::TI, LAYER_SPEC::INTERNAL_BATCH_SIZE, LAYER_SPEC::DIM, BUFFER_SPEC::DYNAMIC_ALLOCATION>;
+            using NOISE_CONTAINER_TYPE = Matrix<NOISE_CONTAINER_SPEC>;
+            NOISE_CONTAINER_TYPE noise;
+        };
+        template<typename T_TYPE_POLICY, typename T_TI, typename T_PARAMETERS = DefaultParameters<T_TYPE_POLICY>>
+        struct Configuration {
+            using TYPE_POLICY = T_TYPE_POLICY;
+            using TI = T_TI;
+            using PARAMETERS = T_PARAMETERS;
+        };
+        template <typename T_CONFIG, typename T_CAPABILITY, typename T_INPUT_SHAPE>
+        struct Specification: T_CAPABILITY, T_CONFIG{
+            using CONFIG = T_CONFIG;
+            using CAPABILITY = T_CAPABILITY;
+            using INPUT_SHAPE = T_INPUT_SHAPE;
+            using TYPE_POLICY = typename CONFIG::TYPE_POLICY;
+            using TI = typename CONFIG::TI;
+            static constexpr TI DIM = get_last(INPUT_SHAPE{});
+            static constexpr TI INPUT_DIM = DIM;
+            static constexpr TI OUTPUT_DIM = DIM;
+            template <typename NEW_INPUT_SHAPE>
+            using OUTPUT_SHAPE_FACTORY = tensor::Replace<NEW_INPUT_SHAPE, OUTPUT_DIM, length(NEW_INPUT_SHAPE{})-1>;
+            using OUTPUT_SHAPE = OUTPUT_SHAPE_FACTORY<INPUT_SHAPE>;
+            static constexpr TI INTERNAL_BATCH_SIZE = get<0>(tensor::CumulativeProduct<tensor::PopBack<INPUT_SHAPE>>{}); // Since the Dense layer is based on Matrices (2D Tensors) the dense layer operation is broadcasted over the leading dimensions. Hence, the actual batch size is the product of all leading dimensions, excluding the last one (containing the features). Since rl_tools::matrix_view is used for zero-cost conversion the INTERNAL_BATCH_SIZE accounts for all leading dimensions.
+            static constexpr TI NUM_WEIGHTS = 0;
+        };
+        template <typename T_SPEC>
+        struct LayerForward{
+            using SPEC = T_SPEC;
+            using TYPE_POLICY = typename SPEC::TYPE_POLICY;
+            using TI = typename SPEC::TI;
+            static constexpr TI DIM = SPEC::DIM;
+            static constexpr TI INPUT_DIM = SPEC::INPUT_DIM;
+            static constexpr TI OUTPUT_DIM = SPEC::OUTPUT_DIM;
+            static constexpr TI INTERNAL_BATCH_SIZE = SPEC::INTERNAL_BATCH_SIZE;
+            using INPUT_SHAPE = typename SPEC::INPUT_SHAPE;
+            using OUTPUT_SHAPE = typename SPEC::OUTPUT_SHAPE;
+            template <typename NEW_INPUT_SHAPE>
+            using OUTPUT_SHAPE_FACTORY = typename SPEC::template OUTPUT_SHAPE_FACTORY<NEW_INPUT_SHAPE>;
+
+            using T = typename TYPE_POLICY::DEFAULT;
+
+            T std = SPEC::PARAMETERS::STD;
+
+            template<bool DYNAMIC_ALLOCATION=true>
+            using State = td3_sampling::State;
+            template<bool DYNAMIC_ALLOCATION=true>
+            using Buffer = td3_sampling::Buffer<td3_sampling::BufferSpecification<TI, SPEC, DYNAMIC_ALLOCATION>>;
+        };
+        template<typename SPEC>
+        struct LayerBackward: public LayerForward<SPEC> {
+            using ACTIVATION_T = typename SPEC::TYPE_POLICY::template GET<numeric_types::categories::Activation>;
+            using PRE_ACTIVATIONS_CONTAINER_SPEC = matrix::Specification<ACTIVATION_T, typename SPEC::TI, SPEC::INTERNAL_BATCH_SIZE, SPEC::DIM, SPEC::DYNAMIC_ALLOCATION>;
+            using PRE_ACTIVATIONS_CONTAINER_TYPE = Matrix<PRE_ACTIVATIONS_CONTAINER_SPEC>;
+            PRE_ACTIVATIONS_CONTAINER_TYPE pre_clip;
+        };
+        template<typename SPEC>
+        struct LayerGradient: public LayerBackward<SPEC> {
+            using T = typename SPEC::TYPE_POLICY::DEFAULT;
+            using OUTPUT_CONTAINER_SPEC = matrix::Specification<T, typename SPEC::TI, SPEC::INTERNAL_BATCH_SIZE, SPEC::DIM, SPEC::DYNAMIC_ALLOCATION>;
+            using OUTPUT_CONTAINER_TYPE = Matrix<OUTPUT_CONTAINER_SPEC>;
+            OUTPUT_CONTAINER_TYPE output;
+        };
+
+        template<typename CONFIG, typename CAPABILITY, typename INPUT_SHAPE>
+        using Layer =
+        typename utils::typing::conditional_t<CAPABILITY::TAG == nn::LayerCapability::Forward,
+                LayerForward<layers::td3_sampling::Specification<CONFIG, CAPABILITY, INPUT_SHAPE>>,
+        typename utils::typing::conditional_t<CAPABILITY::TAG == nn::LayerCapability::Backward,
+                LayerBackward<layers::td3_sampling::Specification<CONFIG, CAPABILITY, INPUT_SHAPE>>,
+        typename utils::typing::conditional_t<CAPABILITY::TAG == nn::LayerCapability::Gradient,
+                LayerGradient<layers::td3_sampling::Specification<CONFIG, CAPABILITY, INPUT_SHAPE>>, void>>>;
+
+        template <typename CONFIG>
+        struct BindConfiguration{
+            template <typename CAPABILITY, typename INPUT_SHAPE>
+            using Layer = nn::layers::td3_sampling::Layer<CONFIG, CAPABILITY, INPUT_SHAPE>;
+        };
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif

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

@@ -0,0 +1,272 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_LAYERS_TD3_SAMPLING_OPERATIONS_GENERIC_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_LAYERS_TD3_SAMPLING_OPERATIONS_GENERIC_H
+#include "../../../nn/activation_functions.h"
+#include "../../../utils/generic/typing.h"
+#include "../../../containers/matrix/matrix.h"
+
+
+#include "layer.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::layers::td3_sampling::LayerForward<SPEC>& layer){ }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::layers::td3_sampling::LayerForward<SPEC>& layer){ }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::layers::td3_sampling::LayerBackward<SPEC>& layer){
+        malloc(device, static_cast<nn::layers::td3_sampling::LayerForward<SPEC>&>(layer));
+        malloc(device, layer.pre_clip);
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::layers::td3_sampling::LayerBackward<SPEC>& layer){
+        free(device, static_cast<nn::layers::td3_sampling::LayerForward<SPEC>&>(layer));
+        free(device, layer.pre_clip);
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::layers::td3_sampling::LayerGradient<SPEC>& layer){
+        malloc(device, static_cast<nn::layers::td3_sampling::LayerBackward<SPEC>&>(layer));
+        malloc(device, layer.output);
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::layers::td3_sampling::LayerGradient<SPEC>& layer){
+        free(device, static_cast<nn::layers::td3_sampling::LayerBackward<SPEC>&>(layer));
+        free(device, layer.output);
+    }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::layers::td3_sampling::Buffer<SPEC>& buffer) {
+        malloc(device, buffer.noise);
+    }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::layers::td3_sampling::Buffer<SPEC>& buffer) {
+        free(device, buffer.noise);
+    }
+    template<typename DEVICE>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::layers::td3_sampling::State& state) { } // no-op
+    template <typename SOURCE_DEVICE, typename TARGET_DEVICE>
+    RL_TOOLS_FUNCTION_PLACEMENT void copy(SOURCE_DEVICE& source_device, TARGET_DEVICE& target_device, nn::layers::td3_sampling::State& source, nn::layers::td3_sampling::State& target){}
+    template<typename DEVICE, typename SPEC, typename RNG, typename MODE>
+    RL_TOOLS_FUNCTION_PLACEMENT void reset(DEVICE& device, const nn::layers::td3_sampling::LayerForward<SPEC>& layer, nn::layers::td3_sampling::State& state, RNG&, Mode<MODE> mode = Mode<mode::Default<>>{}) { } // no-op
+    template<typename DEVICE>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::layers::td3_sampling::State& state) { } // no-op
+    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, nn::layers::td3_sampling::Buffer<SOURCE_SPEC>& source, nn::layers::td3_sampling::Buffer<TARGET_SPEC>& target){
+        copy(source_device, target_device, source.noise, target.noise);
+    }
+    template <typename DEVICE, typename SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT void init_weights(DEVICE& device, nn::layers::td3_sampling::LayerForward<SPEC>& layer, RNG& rng){ }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void reset_forward_state(DEVICE& device, rl_tools::nn::layers::td3_sampling::LayerBackward<SPEC>& l) { }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void zero_gradient(DEVICE& device, nn::layers::td3_sampling::LayerGradient<SPEC>& layer) { }
+    template<typename DEVICE, typename SOURCE_SPEC, typename TARGET_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void add_gradient(DEVICE& device, nn::layers::td3_sampling::LayerGradient<SOURCE_SPEC>& source, nn::layers::td3_sampling::LayerGradient<TARGET_SPEC>& target) {}
+    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::layers::td3_sampling::LayerGradient<SOURCE_SPEC>& source, nn::layers::td3_sampling::LayerGradient<TARGET_SPEC>& target) {}
+    template<typename DEVICE, typename SPEC, typename OPTIMIZER>
+    RL_TOOLS_FUNCTION_PLACEMENT void update(DEVICE& device, nn::layers::td3_sampling::LayerGradient<SPEC>& layer, OPTIMIZER& optimizer){ }
+    template<typename DEVICE, typename SPEC, typename OPTIMIZER>
+    RL_TOOLS_FUNCTION_PLACEMENT void _reset_optimizer_state(DEVICE& device, nn::layers::td3_sampling::LayerGradient<SPEC>& layer, OPTIMIZER& optimizer) { }
+
+    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::layers::td3_sampling::LayerForward<SOURCE_SPEC>& source, nn::layers::td3_sampling::LayerForward<TARGET_SPEC>& target){ }
+    template<typename DEVICE, typename SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT void sample(DEVICE& device, nn::layers::td3_sampling::Buffer<SPEC>& buffer, RNG& rng) {
+        randn(device, buffer.noise, rng);
+    }
+    template <bool SET_PRE_CLIPPING, typename DEVICE, typename SPEC, typename INPUT_SPEC, typename PRE_CLIP_SPEC, typename OUTPUT_SPEC, typename BUFFER_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void evaluate_per_sample(DEVICE& device, const nn::layers::td3_sampling::LayerForward<SPEC>& layer, const Matrix<INPUT_SPEC>& input, Matrix<PRE_CLIP_SPEC>& pre_clip, Matrix<OUTPUT_SPEC>& output, nn::layers::td3_sampling::Buffer<BUFFER_SPEC>& buffer, RNG& rng, typename DEVICE::index_t row_i, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        using TI = typename DEVICE::index_t;
+        using T = typename SPEC::TYPE_POLICY::DEFAULT;
+        using PARAMETERS = typename SPEC::PARAMETERS;
+        for(TI col_i = 0; col_i < SPEC::DIM; col_i++){
+            T mean = get(input, row_i, col_i);
+            add_scalar(device, device.logger, "actor_eval_std", layer.std, 100);
+            T noise;
+            if constexpr(mode::is<MODE, nn::layers::td3_sampling::mode::ExternalNoise>){
+                noise = get(buffer.noise, row_i, col_i);
+            }
+            else{
+                if constexpr(mode::is<MODE, mode::Default> || mode::is<MODE, mode::Rollout>){
+                    noise = random::normal_distribution::sample(device.random, (T)0, (T)1, rng);
+                }
+                else{
+                    if constexpr(mode::is<MODE, mode::Evaluation>){
+                        noise = 0;
+                    }
+                    else{
+                        noise = 0;
+                        utils::assert_exit(device.logger, false, "Invalid mode");
+                    }
+                }
+            }
+            T sample;
+            if constexpr(mode::is<MODE, mode::Evaluation>){
+                sample = mean;
+            }
+            else{
+                sample = mean + noise * layer.std;
+            }
+
+            if constexpr(SET_PRE_CLIPPING){
+                set(pre_clip, row_i, col_i, sample);
+            }
+
+            T sample_clipped = math::clamp(device.math, sample, (T)-1, (T)1);
+            set(output, row_i, col_i, sample_clipped);
+        }
+    }
+    template <typename DEVICE, typename SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC, typename BUFFER_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void evaluate(DEVICE& device, const nn::layers::td3_sampling::LayerForward<SPEC>& layer, const Matrix<INPUT_SPEC>& input, Matrix<OUTPUT_SPEC>& output, nn::layers::td3_sampling::Buffer<BUFFER_SPEC>& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        static_assert(INPUT_SPEC::COLS == SPEC::DIM);
+        static_assert(OUTPUT_SPEC::COLS == SPEC::DIM);
+        static_assert(INPUT_SPEC::ROWS == OUTPUT_SPEC::ROWS);
+        using TI = typename DEVICE::index_t;
+        for(TI row_i = 0; row_i < INPUT_SPEC::ROWS; row_i++){
+            evaluate_per_sample<false>(device, layer, input, output, output, buffer, rng, row_i, mode); // we can pass output for pre_clipping because we don't set it
+        }
+    }
+    template <typename DEVICE, typename SPEC, typename INPUT_SPEC, typename BUFFER_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void forward(DEVICE& device, nn::layers::td3_sampling::LayerGradient<SPEC>& layer, const Matrix<INPUT_SPEC>& input, nn::layers::td3_sampling::Buffer<BUFFER_SPEC>& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        static_assert(INPUT_SPEC::COLS == SPEC::DIM);
+        static_assert(INPUT_SPEC::ROWS == decltype(layer.output)::SPEC::ROWS);
+        using TI = typename DEVICE::index_t;
+        for(TI row_i = 0; row_i < INPUT_SPEC::ROWS; row_i++){
+            evaluate_per_sample<true>(device, layer, input, layer.pre_clip, layer.output, buffer, rng, row_i, mode);
+        }
+    }
+    template <typename DEVICE, typename SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC, typename BUFFER_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void forward(DEVICE& device, nn::layers::td3_sampling::LayerGradient<SPEC>& layer, const Matrix<INPUT_SPEC>& input, Matrix<OUTPUT_SPEC>& output, nn::layers::td3_sampling::Buffer<BUFFER_SPEC>& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        forward(device, layer, input, buffer, rng, mode);
+        copy(device, device, layer.output, output);
+    }
+    template<typename DEVICE, typename SPEC, typename D_OUTPUT_SPEC, typename D_INPUT_SPEC, typename BUFFER_SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void backward_input(DEVICE& device, const nn::layers::td3_sampling::LayerBackward<SPEC>& layer, const Matrix<D_OUTPUT_SPEC>& d_output, Matrix<D_INPUT_SPEC>& d_input, nn::layers::td3_sampling::Buffer<BUFFER_SPEC>&, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        utils::assert_exit(device, false, "Not implemented");
+    }
+    template<typename DEVICE, typename SPEC, typename INPUT_SPEC, typename D_OUTPUT_SPEC, typename BUFFER_SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void backward(DEVICE& device, nn::layers::td3_sampling::LayerGradient<SPEC>& layer, const Matrix<INPUT_SPEC>& input, Matrix<D_OUTPUT_SPEC>& d_output, nn::layers::td3_sampling::Buffer<BUFFER_SPEC>&, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        utils::assert_exit(device, false, "Not implemented");
+    }
+    template<typename DEVICE, typename SPEC, typename INPUT_SPEC, typename D_OUTPUT_SPEC, typename D_INPUT_SPEC, typename BUFFER_SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void backward_full_per_sample(DEVICE& device, nn::layers::td3_sampling::LayerGradient<SPEC>& layer, const Matrix<INPUT_SPEC>& input, Matrix<D_OUTPUT_SPEC>& d_output, Matrix<D_INPUT_SPEC>& d_input, nn::layers::td3_sampling::Buffer<BUFFER_SPEC>&, typename DEVICE::index_t batch_i, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        using T = typename SPEC::TYPE_POLICY::DEFAULT;
+        using TI = typename DEVICE::index_t;
+        constexpr TI ACTION_DIM = SPEC::DIM;
+        for(TI action_i = 0; action_i < ACTION_DIM; action_i++){
+            T action = get(layer.output, batch_i, action_i);
+            T d_output_value = get(d_output, batch_i, action_i);
+            T pre_clip = get(layer.pre_clip, batch_i, action_i);
+            T d_input_value = d_output_value;
+            if(pre_clip < -1 || pre_clip > 1){
+                d_input_value = 0;
+            }
+            set(d_input, batch_i, action_i, d_input_value);
+        }
+    }
+    template<typename DEVICE, typename SPEC, typename INPUT_SPEC, typename D_OUTPUT_SPEC, typename D_INPUT_SPEC, typename BUFFER_SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void backward_full(DEVICE& device, nn::layers::td3_sampling::LayerGradient<SPEC>& layer, const Matrix<INPUT_SPEC>& input, Matrix<D_OUTPUT_SPEC>& d_output, Matrix<D_INPUT_SPEC>& d_input, nn::layers::td3_sampling::Buffer<BUFFER_SPEC>& buffer, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        using TI = typename DEVICE::index_t;
+        using LAYER = nn::layers::td3_sampling::LayerGradient<SPEC>;
+        constexpr TI INTERNAL_BATCH_SIZE = LAYER::INTERNAL_BATCH_SIZE;
+        for(TI batch_i = 0; batch_i < INTERNAL_BATCH_SIZE; batch_i++){
+            backward_full_per_sample(device, layer, input, d_output, d_input, buffer, batch_i, mode);
+        }
+    }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT auto output(DEVICE& device, nn::layers::td3_sampling::LayerGradient<SPEC>& l){
+        auto tensor_flat = to_tensor(device, l.output);
+        auto tensor = view_memory<typename SPEC::OUTPUT_SHAPE>(device, tensor_flat);
+        return tensor;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+// Tensor proxies
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template<typename DEVICE, typename LAYER_SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC, typename BUFFER_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void evaluate(DEVICE& device, const nn::layers::td3_sampling::LayerForward<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<OUTPUT_SPEC>& output, nn::layers::td3_sampling::Buffer<BUFFER_SPEC>& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        auto matrix_view_input = matrix_view(device, input);
+        auto matrix_view_output = matrix_view(device, output);
+        evaluate(device, layer, matrix_view_input, matrix_view_output, buffer, rng, mode);
+    }
+    template<typename DEVICE, typename LAYER_SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC, typename BUFFER_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void evaluate_step(DEVICE& device, const nn::layers::td3_sampling::LayerForward<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, nn::layers::td3_sampling::State& state, Tensor<OUTPUT_SPEC>& output, nn::layers::td3_sampling::Buffer<BUFFER_SPEC>& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        auto matrix_view_input = matrix_view(device, input);
+        auto matrix_view_output = matrix_view(device, output);
+        evaluate(device, layer, matrix_view_input, matrix_view_output, buffer, rng, mode);
+    }
+    template<typename DEVICE, typename LAYER_SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC, typename BUFFER_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void forward(DEVICE& device, nn::layers::td3_sampling::LayerBackward<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<OUTPUT_SPEC>& output, nn::layers::td3_sampling::Buffer<BUFFER_SPEC>& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        auto matrix_view_input = matrix_view(device, input);
+        auto matrix_view_output = matrix_view(device, output);
+        forward(device, layer, matrix_view_input, buffer, rng, mode);
+        copy(device, device, layer.output, matrix_view_output);
+    }
+    template<typename DEVICE, typename LAYER_SPEC, typename INPUT_SPEC, typename BUFFER_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void forward(DEVICE& device, nn::layers::td3_sampling::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, nn::layers::td3_sampling::Buffer<BUFFER_SPEC>& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        auto matrix_view_input = matrix_view(device, input);
+        forward(device, layer, matrix_view_input, buffer, rng, mode);
+    }
+    template<typename DEVICE, typename LAYER_SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC, typename BUFFER_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void forward(DEVICE& device, nn::layers::td3_sampling::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<OUTPUT_SPEC>& output, nn::layers::td3_sampling::Buffer<BUFFER_SPEC>& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        auto matrix_view_input = matrix_view(device, input);
+        auto matrix_view_output = matrix_view(device, output);
+        forward(device, layer, matrix_view_input, matrix_view_output, buffer, rng, mode);
+    }
+    template<typename DEVICE, typename LAYER_SPEC, typename D_OUTPUT_SPEC, typename D_INPUT_SPEC, typename BUFFER_SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void backward_input(DEVICE& device, const nn::layers::td3_sampling::LayerBackward<LAYER_SPEC>& layer, const Tensor<D_OUTPUT_SPEC>& d_output, Tensor<D_INPUT_SPEC>& d_input, nn::layers::td3_sampling::Buffer<BUFFER_SPEC>& buffer, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        auto matrix_view_d_output = matrix_view(device, d_output);
+        auto matrix_view_d_input = matrix_view(device, d_input);
+        backward_input(device, layer, matrix_view_d_output, matrix_view_d_input, buffer, mode);
+    }
+
+    template<typename DEVICE, typename LAYER_SPEC, typename INPUT_SPEC, typename D_OUTPUT_SPEC, typename BUFFER_SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void backward(DEVICE& device, nn::layers::td3_sampling::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<D_OUTPUT_SPEC>& d_output, nn::layers::td3_sampling::Buffer<BUFFER_SPEC>& buffer, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        auto matrix_view_input = matrix_view(device, input);
+        auto matrix_view_d_output = matrix_view(device, d_output);
+        backward(device, layer, matrix_view_input, matrix_view_d_output, buffer, mode);
+    }
+
+    template<typename DEVICE, typename LAYER_SPEC, typename INPUT_SPEC, typename D_OUTPUT_SPEC, typename D_INPUT_SPEC, typename BUFFER_SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void backward_full(DEVICE& device, nn::layers::td3_sampling::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<D_OUTPUT_SPEC>& d_output, Tensor<D_INPUT_SPEC>& d_input, nn::layers::td3_sampling::Buffer<BUFFER_SPEC>& buffer, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        auto matrix_view_input = matrix_view(device, input);
+        auto matrix_view_d_output = matrix_view(device, d_output);
+        auto matrix_view_d_input = matrix_view(device, d_input);
+        backward_full(device, layer, matrix_view_input, matrix_view_d_output, matrix_view_d_input, buffer, mode);
+    }
+    template <typename DEVICE, typename SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT bool is_nan(DEVICE& device, const rl_tools::nn::layers::td3_sampling::LayerForward<SPEC>& l, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        return false;
+    }
+    template <typename DEVICE, typename SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT bool is_nan(DEVICE& device, const rl_tools::nn::layers::td3_sampling::LayerBackward<SPEC>& l, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        bool upstream_nan = is_nan(device, static_cast<const rl_tools::nn::layers::td3_sampling::LayerForward<SPEC>&>(l), mode);
+        if constexpr(mode::is<MODE, nn::parameters::mode::ParametersOnly>){
+            return upstream_nan;
+        }
+        return upstream_nan || is_nan(device, l.pre_squashing, mode) || is_nan(device, l.noise, mode);
+    }
+    template <typename DEVICE, typename SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT bool is_nan(DEVICE& device, const rl_tools::nn::layers::td3_sampling::LayerGradient<SPEC>& l, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        bool upstream_nan = is_nan(device, static_cast<const rl_tools::nn::layers::td3_sampling::LayerBackward<SPEC>&>(l), mode);
+        upstream_nan =  upstream_nan || is_nan(device, l.log_alpha, mode);
+        if constexpr(mode::is<MODE, nn::parameters::mode::ParametersOnly>){
+            return upstream_nan;
+        }
+        return upstream_nan || is_nan(device, l.log_probabilities, mode) || is_nan(device, l.output, mode);
+    }
+    template<typename DEVICE, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT bool is_nan(DEVICE& device, nn::layers::td3_sampling::State& state, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        return false;
+    }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT auto gradient_norm(DEVICE& device, const nn::layers::td3_sampling::LayerGradient<SPEC>& layer) {
+        return 0;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/nn/layers/td3_sampling/persist.h

@@ -0,0 +1,17 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_LAYERS_TD3_SAMPLING_PERSIST_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_LAYERS_TD3_SAMPLING_PERSIST_H
+
+#include "layer.h"
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools {
+    template<typename DEVICE, typename SPEC, typename GROUP>
+    void save(DEVICE& device, nn::layers::td3_sampling::LayerForward<SPEC>& layer, GROUP& group) { }
+    template<typename DEVICE, typename SPEC, typename GROUP>
+    bool load(DEVICE& device, nn::layers::td3_sampling::LayerForward<SPEC>& layer, GROUP& group) {
+        return true;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/nn/layers/td3_sampling/persist_code.h

@@ -0,0 +1,164 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_LAYERS_TD3_SAMPLING_PERSIST_CODE_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_LAYERS_TD3_SAMPLING_PERSIST_CODE_H
+#include "layer.h"
+#include "../../../containers/matrix/persist_code.h"
+#include <sstream>
+#include "../../../persist/code.h"
+#include "../../../containers/matrix/persist_code.h"
+#include "../../../nn/capability/persist_code.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools {
+    namespace nn::layers::td3_sampling::persist_code{
+        template<typename DEVICE, typename SPEC>
+            rl_tools::persist::Code finish(DEVICE& device, nn::layers::td3_sampling::LayerForward<SPEC> &layer, std::string name, rl_tools::persist::Code input, bool const_declaration=true, typename DEVICE::index_t indent=0){
+            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();
+            using TI = typename DEVICE::index_t;
+            std::stringstream ss, ss_header;
+            ss_header << input.header;
+            ss_header << "#include <rl_tools/nn/layers/td3_sampling/layer.h>\n";
+            ss << input.body;
+            std::string T_string = containers::persist::get_type_string<typename SPEC::TYPE_POLICY::DEFAULT>();
+            std::string TI_string = containers::persist::get_type_string<typename SPEC::TI>();
+            ss << ind << "namespace " << name << " {\n";
+            std::string T_parameter_string = containers::persist::get_type_string<typename SPEC::TYPE_POLICY::DEFAULT>();
+            ss << ind << "    using PARAMETERS = " << "struct PARAMETERS{";
+            ss << ind << "        static constexpr " << T_parameter_string << " STD = " << SPEC::PARAMETERS::STD << ";\n";
+            ss << ind << "    };\n";
+            ss << ind << "    using CONFIG = " << "RL_TOOLS""_NAMESPACE_WRAPPER ::rl_tools::nn::layers::td3_sampling::Configuration<";
+            ss << T_string << ", ";
+            ss << TI_string << ", ";
+            ss << "PARAMETERS";
+            ss << ">; \n";;
+            ss << ind << "    " << "using TEMPLATE = RL_TOOLS""_NAMESPACE_WRAPPER ::rl_tools::nn::layers::td3_sampling::BindConfiguration<CONFIG>;" << "\n";
+            ss << ind << "    " << "using INPUT_SHAPE = RL_TOOLS""_NAMESPACE_WRAPPER ::rl_tools::tensor::Shape<" << TI_string << ", " << get<0>(typename SPEC::INPUT_SHAPE{}) << ", " << get<1>(typename SPEC::INPUT_SHAPE{}) << ", " << get<2>(typename SPEC::INPUT_SHAPE{}) << ">;\n";
+            ss << ind << "    " << "using CAPABILITY = " << to_string(typename SPEC::CAPABILITY{}) << ";" << "\n";
+            ss << ind << "    " << "using TYPE = RL_TOOLS""_NAMESPACE_WRAPPER ::rl_tools::nn::layers::td3_sampling::Layer<CONFIG, CAPABILITY, INPUT_SHAPE>;" << "\n";
+            std::string initializer_list;
+            if constexpr(SPEC::CAPABILITY::TAG == nn::LayerCapability::Forward){
+                initializer_list = "{}";
+            }
+            else{
+                if constexpr(SPEC::CAPABILITY::TAG == nn::LayerCapability::Backward){
+                    initializer_list = "{{}, pre_squashing::container, noise::container}";
+                }
+                else{
+                    if constexpr(SPEC::CAPABILITY::TAG == nn::LayerCapability::Gradient){
+                        initializer_list = "{{{}, pre_squashing::container, noise::container}, log_probabilities::container, output::container, log_alpha::parameters}";
+                    }
+                    else{
+                        utils::assert_exit(device, false, "Unknown capability");
+                    }
+                }
+            }
+            ss << ind << "    " << (const_declaration ? "constexpr " : "") << "TYPE module = " << initializer_list << ";\n";
+            ss << ind << "    " << "template <typename T_TYPE = TYPE>" << "\n";
+            ss << ind << "    " << (const_declaration ? "constexpr " : "") << "T_TYPE factory = " << initializer_list << ";" << "\n";
+            ss << ind << "    " << "template <typename T_TYPE = TYPE>" << "\n";
+            ss << ind << "    " << (const_declaration ? "constexpr " : "") << "T_TYPE factory_function(){return T_TYPE" << initializer_list << ";" << "}\n";
+            ss << ind << "}\n";
+
+
+            return {ss_header.str(), ss.str()};
+        }
+    }
+    template<typename DEVICE, typename SPEC>
+    persist::Code save_code_split(DEVICE& device, nn::layers::td3_sampling::LayerForward<SPEC> &layer, std::string name, bool const_declaration=true, typename DEVICE::index_t indent=0, bool finish=true){
+        if(finish){
+            return nn::layers::td3_sampling::persist_code::finish(device, layer, name, {"", ""}, const_declaration, indent);
+        }
+        else{
+            return {"", ""};
+        }
+    }
+    template<typename DEVICE, typename SPEC>
+    persist::Code save_code_split(DEVICE& device, nn::layers::td3_sampling::LayerBackward<SPEC> &layer, std::string name, bool const_declaration=true, typename DEVICE::index_t indent=0, bool finish=true){
+        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();
+        using TI = typename DEVICE::index_t;
+        std::stringstream ss, ss_header;
+        auto previous = save_code_split(device, static_cast<nn::layers::td3_sampling::LayerForward<SPEC>&>(layer), name, const_declaration, indent, false);
+        ss_header << previous.header;
+        ss << previous.body;
+        ss << ind << "namespace " << name << " {\n";
+        {
+            auto pre_activations = save_code_split(device, layer.pre_squashing, "pre_squashing", const_declaration, indent+1);
+            ss_header << pre_activations.header;
+            ss << pre_activations.body;
+        }
+        {
+            auto pre_activations = save_code_split(device, layer.noise, "noise", const_declaration, indent+1);
+            ss_header << pre_activations.header;
+            ss << pre_activations.body;
+        }
+        ss << ind << "}\n";
+        if(finish){
+            return nn::layers::td3_sampling::persist_code::finish(device, layer, name, {ss_header.str(), ss.str()}, const_declaration, indent);
+        }
+        else{
+            return {ss_header.str(), ss.str()};
+        }
+    }
+
+    template<typename DEVICE, typename SPEC>
+    persist::Code save_code_split(DEVICE& device, nn::layers::td3_sampling::LayerGradient<SPEC> &layer, std::string name, bool const_declaration=true, typename DEVICE::index_t indent=0){
+        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();
+        using TI = typename DEVICE::index_t;
+        std::stringstream ss, ss_header;
+        auto previous = save_code_split(device, static_cast<nn::layers::td3_sampling::LayerBackward<SPEC>&>(layer), name, const_declaration, indent, false);
+        ss_header << previous.header;
+        ss << previous.body;
+        ss << ind << "namespace " << name << " {\n";
+        {
+            auto output = save_code_split(device, layer.log_alpha, "log_alpha", const_declaration, indent+1);
+            ss_header << output.header;
+            ss << output.body;
+        }
+        {
+            auto output = save_code_split(device, layer.log_probabilities, "log_probabilities", const_declaration, indent+1);
+            ss_header << output.header;
+            ss << output.body;
+        }
+        {
+            auto output = save_code_split(device, layer.output, "output", const_declaration, indent+1);
+            ss_header << output.header;
+            ss << output.body;
+        }
+        ss << ind << "}\n";
+        return nn::layers::td3_sampling::persist_code::finish(device, layer, name, {ss_header.str(), ss.str()}, const_declaration, indent);
+    }
+
+    template<typename DEVICE, typename SPEC>
+    std::string save_code(DEVICE& device, nn::layers::td3_sampling::LayerForward<SPEC> &layer, std::string name, bool const_declaration=true, typename DEVICE::index_t indent=0){
+        auto code = save_code_split(device, layer, name, const_declaration, indent);
+        return code.header + code.body;
+    }
+    template <typename DEVICE, typename SPEC>
+    std::string nn_analytics(DEVICE& device, nn::layers::td3_sampling::LayerGradient<SPEC>& layer) {
+        std::string data;
+        data += "{";
+        data += "\"pre_clip\": " + json(device, layer.pre_clip) + ", ";
+        data += "\"output\": " + json(device, layer.output);
+        data += "}";
+        return data;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif