hyperray 0.0.1__py3-none-any.whl

hyperdrone/_vendor/rl-tools/include/rl_tools/nn_models/multi_agent_wrapper/model.h

@@ -0,0 +1,148 @@
+#include "../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_MODELS_MULTI_AGENT_WRAPPER_MODEL_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_MODELS_MULTI_AGENT_WRAPPER_MODEL_H
+
+#include "../../nn/nn.h"
+#include "../../nn_models/sequential/model.h"
+#include "../../nn/parameters/parameters.h"
+#include "../../nn/optimizers/sgd/sgd.h"
+#include "../../nn/optimizers/adam/adam.h"
+#include "../../utils/generic/typing.h"
+#include "../../containers/matrix/matrix.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::nn_models::multi_agent_wrapper {
+    template <typename T_TYPE_POLICY, typename T_TI, T_TI T_N_AGENTS, typename T_MODULE>
+    struct Configuration{
+        using TYPE_POLICY = T_TYPE_POLICY;
+        using TI = T_TI;
+        static constexpr TI N_AGENTS = T_N_AGENTS;
+        using MODULE = T_MODULE;
+    };
+
+//    template <typename T_T, typename T_TI, T_TI T_N_AGENTS, typename T_MODULE>
+//    using SpecificationFixedModule = Specification<T_T, T_TI, T_N_AGENTS, BindModule<T_MODULE>::template Module>;
+
+    template <typename T_CONFIG, typename T_CAPABILITY, typename T_INPUT_SHAPE>
+    struct Specification: T_CAPABILITY, T_CONFIG{
+        using TYPE_POLICY = typename T_CONFIG::TYPE_POLICY;
+        using TI = typename T_CONFIG::TI;
+        using CONFIG = T_CONFIG;
+        using CAPABILITY = T_CAPABILITY;
+
+        using PARAMETER_TYPE = typename CAPABILITY::PARAMETER_TYPE;
+
+        static constexpr TI INPUT_DIM = get_last(T_INPUT_SHAPE{});
+        static constexpr TI BATCH_AXIS = length(T_INPUT_SHAPE{})-2;
+        static constexpr TI BATCH_SIZE = get<BATCH_AXIS>(T_INPUT_SHAPE{});
+        static_assert(INPUT_DIM % CONFIG::N_AGENTS == 0, "The input dimension must be divisible by the number of agents");
+        static constexpr TI PER_AGENT_INPUT_DIM = INPUT_DIM / CONFIG::N_AGENTS;
+        static constexpr TI INTERNAL_BATCH_SIZE = get<0>(tensor::CumulativeProduct<tensor::PopBack<T_INPUT_SHAPE>>{}) * CONFIG::N_AGENTS; // 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.
+        using INPUT_SHAPE = T_INPUT_SHAPE;
+        using INTERNAL_INPUT_SHAPE_TEMP = tensor::Replace<T_INPUT_SHAPE, PER_AGENT_INPUT_DIM, length(T_INPUT_SHAPE{})-1>;
+        using INTERNAL_INPUT_SHAPE = tensor::Replace<INTERNAL_INPUT_SHAPE_TEMP, INTERNAL_BATCH_SIZE, BATCH_AXIS>;
+        using MODEL = nn_models::sequential::Build<CAPABILITY, typename CONFIG::MODULE, INTERNAL_INPUT_SHAPE>;
+        static constexpr TI PER_AGENT_OUTPUT_DIM = get_last(typename MODEL::OUTPUT_SHAPE{});
+        static constexpr TI OUTPUT_DIM = PER_AGENT_OUTPUT_DIM * CONFIG::N_AGENTS;
+        using OUTPUT_SHAPE = tensor::Replace<INPUT_SHAPE, OUTPUT_DIM, length(typename MODEL::INPUT_SHAPE{})-1>;
+    };
+
+    template <typename T_SPEC, bool T_DYNAMIC_ALLOCATION=true>
+    struct ModuleStateSpecification{
+        using SPEC = T_SPEC;
+        static constexpr bool DYNAMIC_ALLOCATION = T_DYNAMIC_ALLOCATION;
+    };
+
+    template<typename T_SPEC>
+    struct ModuleState{
+        using STATE_SPEC = T_SPEC;
+        using SPEC = typename STATE_SPEC::SPEC;
+
+        using INNER_STATE = typename SPEC::MODEL::template State<SPEC::DYNAMIC_ALLOCATION>;
+        INNER_STATE inner_state;
+    };
+
+    template<typename T_SPEC, typename T_MODEL, bool T_DYNAMIC_ALLOCATION=true>
+    struct ModuleBuffersSpecification{
+        using SPEC = T_SPEC;
+        using TI = typename SPEC::TI;
+        static constexpr TI BATCH_SIZE = SPEC::BATCH_SIZE;
+        using MODEL = T_MODEL;
+        static constexpr bool DYNAMIC_ALLOCATION = T_DYNAMIC_ALLOCATION;
+    };
+
+    template<typename T_BUFFER_SPEC>
+    struct ModuleBuffer{
+        using BUFFER_SPEC = T_BUFFER_SPEC;
+        using SPEC = typename BUFFER_SPEC::SPEC;
+        using TYPE_POLICY = typename SPEC::TYPE_POLICY;
+        using TI = typename SPEC::TI;
+        static constexpr TI BATCH_SIZE = T_BUFFER_SPEC::BATCH_SIZE;
+        static constexpr TI INNER_BATCH_SIZE = BATCH_SIZE * SPEC::N_AGENTS;
+
+        using T_INPUT = typename TYPE_POLICY::template GET<numeric_types::categories::Input>;
+        using T_OUTPUT = typename TYPE_POLICY::template GET<numeric_types::categories::Accumulator>;
+
+        using INPUT_BUFFER_SPEC = tensor::Specification<T_INPUT, TI, tensor::Shape<TI, 1, BATCH_SIZE, SPEC::INPUT_DIM>, BUFFER_SPEC::DYNAMIC_ALLOCATION>;
+        using INPUT_BUFFER_TYPE = Tensor<INPUT_BUFFER_SPEC>;
+        using OUTPUT_BUFFER_SPEC = tensor::Specification<T_OUTPUT, TI, tensor::Shape<TI, 1, BATCH_SIZE, SPEC::OUTPUT_DIM>, BUFFER_SPEC::DYNAMIC_ALLOCATION>;
+        using OUTPUT_BUFFER_TYPE = Tensor<OUTPUT_BUFFER_SPEC>;
+        INPUT_BUFFER_TYPE input, d_input;
+        OUTPUT_BUFFER_TYPE output;
+
+        typename BUFFER_SPEC::MODEL::template Buffer<BUFFER_SPEC::DYNAMIC_ALLOCATION> buffer;
+    };
+
+    template<typename T_SPEC>
+    struct ModuleForward{
+        using SPEC = T_SPEC;
+        using TYPE_POLICY = typename SPEC::TYPE_POLICY;
+        using TI = typename SPEC::TI;
+        using MODEL = typename SPEC::MODEL;
+        static constexpr TI N_AGENTS = SPEC::N_AGENTS;
+        using INPUT_SHAPE = typename SPEC::INPUT_SHAPE;
+        using OUTPUT_SHAPE = typename SPEC::OUTPUT_SHAPE;
+        MODEL content;
+        template <bool DYNAMIC_ALLOCATION=true>
+        using State = ModuleState<ModuleStateSpecification<SPEC,  DYNAMIC_ALLOCATION>>;
+        template<bool DYNAMIC_ALLOCATION=true>
+        using Buffer = ModuleBuffer<ModuleBuffersSpecification<SPEC, MODEL, DYNAMIC_ALLOCATION>>;
+    };
+
+    template<typename SPEC>
+    struct ModuleBackward: public ModuleForward<SPEC>{
+//        static constexpr typename SPEC::TI BATCH_SIZE = SPEC::BATCH_SIZE;
+    };
+    template<typename T_SPEC>
+    struct ModuleGradient: public ModuleBackward<T_SPEC>{
+        using TI = typename T_SPEC::TI;
+    };
+
+    template <typename T_CAPABILITY, typename T_SPEC>
+    using UpgradeCapabilityBatchSize = typename T_CAPABILITY::template CHANGE_BATCH_SIZE<T_CAPABILITY::BATCH_SIZE*T_SPEC::N_AGENTS>;
+
+    template<typename CONFIG, typename CAPABILITY, typename INPUT_SHAPE>
+    using Module =
+            typename utils::typing::conditional_t<CAPABILITY::TAG == nn::LayerCapability::Forward, ModuleForward<Specification<CONFIG, CAPABILITY, INPUT_SHAPE>>,
+            typename utils::typing::conditional_t<CAPABILITY::TAG == nn::LayerCapability::Backward, ModuleBackward<Specification<CONFIG, CAPABILITY, INPUT_SHAPE>>,
+            typename utils::typing::conditional_t<CAPABILITY::TAG == nn::LayerCapability::Gradient, ModuleGradient<Specification<CONFIG, CAPABILITY, INPUT_SHAPE>>, void>>>;
+
+
+    template <typename CAPABILITY, typename CONFIG, typename INPUT_SHAPE>
+    struct Build: Module<CONFIG, CAPABILITY, INPUT_SHAPE>{
+        template <typename TI, TI BATCH_SIZE>
+        struct CHANGE_BATCH_SIZE_IMPL{
+            using NEW_INPUT_SHAPE = tensor::Replace<INPUT_SHAPE, BATCH_SIZE, 1>;
+            using CHANGE_BATCH_SIZE = Build<CAPABILITY, CONFIG, NEW_INPUT_SHAPE>;
+        };
+        template <typename TI, TI BATCH_SIZE>
+        using CHANGE_BATCH_SIZE = typename CHANGE_BATCH_SIZE_IMPL<TI, BATCH_SIZE>::CHANGE_BATCH_SIZE;
+        template <typename NEW_CAPABILITY>
+        using CHANGE_CAPABILITY = Build<NEW_CAPABILITY, CONFIG, INPUT_SHAPE>;
+    };
+
+
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/nn_models/multi_agent_wrapper/operations_generic.h

@@ -0,0 +1,288 @@
+#include "../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_MODELS_MULTI_AGENT_WRAPPER_OPERATIONS_GENERIC_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_MODELS_MULTI_AGENT_WRAPPER_OPERATIONS_GENERIC_H
+
+#include "model.h"
+#include "../../utils/generic/typing.h"
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template <typename DEVICE, typename MODULE_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn_models::multi_agent_wrapper::ModuleForward<MODULE_SPEC>& module){
+        malloc(device, module.content);
+    }
+    template <typename DEVICE, typename MODULE_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn_models::multi_agent_wrapper::ModuleForward<MODULE_SPEC>& module){
+        free(device, module.content);
+    }
+    template <typename DEVICE, typename STATE_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn_models::multi_agent_wrapper::ModuleState<STATE_SPEC>& state){
+        malloc(device, state.inner_state);
+    }
+    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_models::multi_agent_wrapper::ModuleState<SOURCE_SPEC>& source, nn_models::multi_agent_wrapper::ModuleState<TARGET_SPEC>& target){
+        copy(source_device, target_device, source.inner_state, target.inner_state);
+    }
+    template <typename DEVICE, typename MODULE_SPEC, typename STATE_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void reset(DEVICE& device, const nn_models::multi_agent_wrapper::ModuleForward<MODULE_SPEC>& model, nn_models::multi_agent_wrapper::ModuleState<STATE_SPEC>& state, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        reset(device, model.content, state.inner_state, rng, mode);
+    }
+    template <typename DEVICE, typename STATE_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn_models::multi_agent_wrapper::ModuleState<STATE_SPEC>& state){
+        free(device, state.inner_state);
+    }
+    template <typename DEVICE, typename BUFFER_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn_models::multi_agent_wrapper::ModuleBuffer<BUFFER_SPEC>& buffer){
+        malloc(device, buffer.input);
+        malloc(device, buffer.d_input);
+        malloc(device, buffer.output);
+        malloc(device, buffer.buffer);
+    }
+    template <typename DEVICE, typename BUFFER_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn_models::multi_agent_wrapper::ModuleBuffer<BUFFER_SPEC>& buffer){
+        free(device, buffer.input);
+        free(device, buffer.d_input);
+        free(device, buffer.output);
+        free(device, buffer.buffer);
+    }
+    template <typename SOURCE_DEVICE, typename TARGET_DEVICE, typename SOURCE_BUFFER_SPEC, typename TARGET_BUFFER_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void copy(SOURCE_DEVICE& source_device, TARGET_DEVICE& target_device, nn_models::multi_agent_wrapper::ModuleBuffer<SOURCE_BUFFER_SPEC>& source, nn_models::multi_agent_wrapper::ModuleBuffer<TARGET_BUFFER_SPEC>& target){
+        copy(source_device, target_device, source.buffer, target.buffer);
+    }
+    template <typename DEVICE, typename MODULE_SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT void init_weights(DEVICE& device, nn_models::multi_agent_wrapper::ModuleForward<MODULE_SPEC>& module, RNG& rng){
+        init_weights(device, module.content, rng);
+    }
+
+    template <typename MODULE_SPEC> // non-const
+    RL_TOOLS_FUNCTION_PLACEMENT constexpr auto& get_first_layer(nn_models::multi_agent_wrapper::ModuleForward<MODULE_SPEC>& model){
+        return get_first_layer(model.content);
+    }
+    template <typename MODULE_SPEC> // const
+    RL_TOOLS_FUNCTION_PLACEMENT constexpr auto& get_first_layer(const nn_models::multi_agent_wrapper::ModuleForward<MODULE_SPEC>& model){
+        return get_first_layer(model.content);
+    }
+
+    template <typename MODULE_SPEC> // non-const
+    RL_TOOLS_FUNCTION_PLACEMENT constexpr auto& get_last_layer(nn_models::multi_agent_wrapper::ModuleForward<MODULE_SPEC>& model){
+        return get_last_layer(model.content);
+    }
+    template <typename MODULE_SPEC> // const
+    RL_TOOLS_FUNCTION_PLACEMENT constexpr auto& get_last_layer(const nn_models::multi_agent_wrapper::ModuleForward<MODULE_SPEC>& model){
+        return get_last_layer(model.content);
+    }
+
+    template <typename BUFFER_SPEC> // non-const
+    RL_TOOLS_FUNCTION_PLACEMENT constexpr auto& get_last_buffer(nn_models::multi_agent_wrapper::ModuleBuffer<BUFFER_SPEC>& buffer){
+        return get_last_buffer(buffer.buffer);
+    }
+    template <typename BUFFER_SPEC> // const
+    RL_TOOLS_FUNCTION_PLACEMENT constexpr auto& get_last_buffer(const nn_models::multi_agent_wrapper::ModuleBuffer<BUFFER_SPEC>& buffer){
+        return get_last_buffer(buffer.buffer);
+    }
+
+    template <typename DEVICE, typename SPEC> // non-const
+    RL_TOOLS_FUNCTION_PLACEMENT constexpr auto output(DEVICE& device, nn_models::multi_agent_wrapper::ModuleGradient<SPEC>& m){
+        auto output_reshaped = reshape_row_major(device, output(device, m.content), typename SPEC::OUTPUT_SHAPE{});
+        return output_reshaped;
+    }
+    template <typename DEVICE, typename SPEC> // const
+    RL_TOOLS_FUNCTION_PLACEMENT constexpr auto output(DEVICE& device, const nn_models::multi_agent_wrapper::ModuleGradient<SPEC>& m){
+        auto output_reshaped = reshape_row_major(device, output(device, m.content), typename SPEC::OUTPUT_SHAPE{});
+        return output_reshaped;
+    }
+    template<typename DEVICE, typename MODULE_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_models::multi_agent_wrapper::ModuleForward<MODULE_SPEC>& model, const Tensor<INPUT_SPEC>& input, Tensor<OUTPUT_SPEC>& output, nn_models::multi_agent_wrapper::ModuleBuffer<BUFFER_SPEC>& buffers, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        using TI = typename DEVICE::index_t;
+        static_assert(INPUT_SPEC::SHAPE::LENGTH == 3);
+        static_assert(OUTPUT_SPEC::SHAPE::LENGTH == 3);
+        constexpr TI INPUT_DIM = INPUT_SPEC::SHAPE::template GET<INPUT_SPEC::SHAPE::LENGTH-1>;
+        constexpr TI OUTPUT_DIM = OUTPUT_SPEC::SHAPE::template GET<OUTPUT_SPEC::SHAPE::LENGTH-1>;
+        static_assert(INPUT_DIM % MODULE_SPEC::N_AGENTS == 0);
+        static_assert(OUTPUT_DIM % MODULE_SPEC::N_AGENTS == 0);
+        constexpr TI BATCH_AXIS = 1;
+        constexpr TI BATCH_SIZE = INPUT_SPEC::SHAPE::template GET<BATCH_AXIS>;
+
+        auto input_buffer_view = view_range(device, buffers.input, tensor::ViewSpec<BATCH_AXIS, BATCH_SIZE>{});
+        auto output_buffer_view = view_range(device, buffers.output, tensor::ViewSpec<BATCH_AXIS, BATCH_SIZE>{});
+        copy(device, device, input, input_buffer_view);// forgot why this? to make it dense for the reshape?
+        auto input_inner = reshape_row_major(device, input_buffer_view, tensor::Shape<TI, BATCH_SIZE*MODULE_SPEC::N_AGENTS, INPUT_DIM/MODULE_SPEC::N_AGENTS>{});
+        auto output_inner = reshape_row_major(device, output_buffer_view, tensor::Shape<TI, BATCH_SIZE*MODULE_SPEC::N_AGENTS, OUTPUT_DIM/MODULE_SPEC::N_AGENTS>{});
+        evaluate(device, model.content, input_inner, output_inner, buffers.buffer, rng, mode);
+        copy(device, device, output_buffer_view, output);
+    }
+    template<typename DEVICE, typename MODULE_SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC, typename STATE_SPEC, typename BUFFER_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void evaluate_step(DEVICE& device, const nn_models::multi_agent_wrapper::ModuleForward<MODULE_SPEC>& model, const Tensor<INPUT_SPEC>& input, nn_models::multi_agent_wrapper::ModuleState<STATE_SPEC>& state, Tensor<OUTPUT_SPEC>& output, nn_models::multi_agent_wrapper::ModuleBuffer<BUFFER_SPEC>& buffers, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        using TI = typename DEVICE::index_t;
+        static_assert(INPUT_SPEC::SHAPE::LENGTH == 2);
+        static_assert(OUTPUT_SPEC::SHAPE::LENGTH == 2);
+        constexpr TI INPUT_DIM = INPUT_SPEC::SHAPE::template GET<INPUT_SPEC::SHAPE::LENGTH-1>;
+        constexpr TI OUTPUT_DIM = OUTPUT_SPEC::SHAPE::template GET<OUTPUT_SPEC::SHAPE::LENGTH-1>;
+        static_assert(INPUT_DIM % MODULE_SPEC::N_AGENTS == 0);
+        static_assert(OUTPUT_DIM % MODULE_SPEC::N_AGENTS == 0);
+        constexpr TI BATCH_AXIS = 0;
+        constexpr TI BATCH_SIZE = INPUT_SPEC::SHAPE::template GET<BATCH_AXIS>;
+        using TI = typename DEVICE::index_t;
+
+        auto input_buffer_step = view(device, buffers.input, 0);
+        auto output_buffer_step = view(device, buffers.output, 0);
+        auto input_buffer_view = view_range(device, input_buffer_step, tensor::ViewSpec<BATCH_AXIS, BATCH_SIZE>{});
+        auto output_buffer_view = view_range(device, output_buffer_step, tensor::ViewSpec<BATCH_AXIS, BATCH_SIZE>{});
+        copy(device, device, input, input_buffer_view);
+        auto input_inner = reshape_row_major(device, input_buffer_view, tensor::Shape<TI, BATCH_SIZE*MODULE_SPEC::N_AGENTS, INPUT_DIM/MODULE_SPEC::N_AGENTS>{});
+        auto output_inner = reshape_row_major(device, output_buffer_view, tensor::Shape<TI, BATCH_SIZE*MODULE_SPEC::N_AGENTS, OUTPUT_DIM/MODULE_SPEC::N_AGENTS>{});
+        evaluate_step(device, model.content, input_inner, state.inner_state, output_inner, buffers.buffer, rng, mode);
+        copy(device, device, output_buffer_view, output);
+    }
+    template <typename DEVICE, typename MODULE_SPEC, typename INPUT_SPEC, typename BUFFER_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void forward(DEVICE& device, nn_models::multi_agent_wrapper::ModuleGradient<MODULE_SPEC>& module, const Tensor<INPUT_SPEC>& input, nn_models::multi_agent_wrapper::ModuleBuffer<BUFFER_SPEC>& buffers, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        using TI = typename DEVICE::index_t;
+        static_assert(INPUT_SPEC::SHAPE::LENGTH == 3);
+        constexpr TI INPUT_DIM = INPUT_SPEC::SHAPE::template GET<INPUT_SPEC::SHAPE::LENGTH-1>;
+        static_assert(INPUT_DIM % MODULE_SPEC::N_AGENTS == 0);
+        constexpr TI BATCH_AXIS = 1;
+        constexpr TI BATCH_SIZE = INPUT_SPEC::SHAPE::template GET<BATCH_AXIS>;
+        using TI = typename DEVICE::index_t;
+
+        auto input_buffer_view = view_range(device, buffers.input, tensor::ViewSpec<BATCH_AXIS, BATCH_SIZE>{});
+        auto output_buffer_view = view_range(device, buffers.output, tensor::ViewSpec<BATCH_AXIS, BATCH_SIZE>{});
+        copy(device, device, input, input_buffer_view);
+        auto input_inner = reshape_row_major(device, input_buffer_view, tensor::Shape<TI, BATCH_SIZE*MODULE_SPEC::N_AGENTS, INPUT_DIM/MODULE_SPEC::N_AGENTS>{});
+        forward(device, module.content, input_inner, buffers.buffer, rng, mode);
+    }
+    template <typename DEVICE, typename MODULE_SPEC, typename INPUT, typename OUTPUT, typename BUFFER_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void forward(DEVICE& device, nn_models::multi_agent_wrapper::ModuleGradient<MODULE_SPEC>& module, INPUT& input, OUTPUT& output, nn_models::multi_agent_wrapper::ModuleBuffer<BUFFER_SPEC>& buffers, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        forward(device, module, input, buffers, rng, mode);
+        auto module_output = rl_tools::output(device, module);
+        copy(device, device, module_output, output);
+    }
+    template <typename DEVICE, typename MODULE_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void zero_gradient(DEVICE& device, nn_models::multi_agent_wrapper::ModuleGradient<MODULE_SPEC>& module){
+        zero_gradient(device, module.content);
+    }
+    template<typename DEVICE, typename SOURCE_SPEC, typename TARGET_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void add_gradient(DEVICE& device, nn_models::multi_agent_wrapper::ModuleGradient<SOURCE_SPEC>& source, nn_models::multi_agent_wrapper::ModuleGradient<TARGET_SPEC>& target) {
+        add_gradient(device, source.content, target.content);
+    }
+    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_models::multi_agent_wrapper::ModuleGradient<SOURCE_SPEC>& source, nn_models::multi_agent_wrapper::ModuleGradient<TARGET_SPEC>& target) {
+        copy_gradient(source_device, target_device, source.content, target.content);
+    }
+    template<typename DEVICE, typename SPEC, typename OPTIMIZER>
+    RL_TOOLS_FUNCTION_PLACEMENT void _reset_optimizer_state(DEVICE& device, nn_models::multi_agent_wrapper::ModuleGradient<SPEC>& module, OPTIMIZER& optimizer) {
+        _reset_optimizer_state(device, module.content, optimizer);
+    }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void reset_forward_state(DEVICE& device, nn_models::multi_agent_wrapper::ModuleGradient<SPEC>& module) {
+        reset_forward_state(device, module.content);
+    }
+    template<typename DEVICE, typename MODULE_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_models::multi_agent_wrapper::ModuleGradient<MODULE_SPEC>& model, const Tensor<INPUT_SPEC>& input, Tensor<D_OUTPUT_SPEC>& d_output, Tensor<D_INPUT_SPEC>& d_input, nn_models::multi_agent_wrapper::ModuleBuffer<BUFFER_SPEC> buffers, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        using TI = typename DEVICE::index_t;
+        static_assert(INPUT_SPEC::SHAPE::LENGTH == 3);
+        static_assert(D_INPUT_SPEC::SHAPE::LENGTH == 3);
+        static_assert(D_OUTPUT_SPEC::SHAPE::LENGTH == 3);
+        constexpr TI INPUT_DIM = INPUT_SPEC::SHAPE::template GET<INPUT_SPEC::SHAPE::LENGTH-1>;
+        static_assert(INPUT_DIM == D_INPUT_SPEC::SHAPE::template GET<D_INPUT_SPEC::SHAPE::LENGTH-1>);
+        constexpr TI OUTPUT_DIM = D_OUTPUT_SPEC::SHAPE::template GET<D_OUTPUT_SPEC::SHAPE::LENGTH-1>;
+        static_assert(INPUT_DIM % MODULE_SPEC::N_AGENTS == 0);
+        static_assert(OUTPUT_DIM % MODULE_SPEC::N_AGENTS == 0);
+        constexpr TI BATCH_AXIS = 1;
+        constexpr TI BATCH_SIZE = INPUT_SPEC::SHAPE::template GET<BATCH_AXIS>;
+        static_assert(BATCH_SIZE == D_INPUT_SPEC::SHAPE::template GET<BATCH_AXIS>);
+        static_assert(BATCH_SIZE == D_OUTPUT_SPEC::SHAPE::template GET<BATCH_AXIS>);
+        using TI = typename DEVICE::index_t;
+
+        auto input_buffer_view = view_range(device, buffers.input, tensor::ViewSpec<BATCH_AXIS, BATCH_SIZE>{});
+        auto d_input_buffer_view = view_range(device, buffers.d_input, tensor::ViewSpec<BATCH_AXIS, BATCH_SIZE>{});
+        auto d_output_buffer_view = view_range(device, buffers.output, tensor::ViewSpec<BATCH_AXIS, BATCH_SIZE>{});
+        copy(device, device, input, input_buffer_view);
+        copy(device, device, d_output, d_output_buffer_view);
+        auto input_inner = reshape_row_major(device, input_buffer_view, tensor::Shape<TI, BATCH_SIZE*MODULE_SPEC::N_AGENTS, INPUT_DIM/MODULE_SPEC::N_AGENTS>{});
+        auto d_input_inner = reshape_row_major(device, d_input_buffer_view, tensor::Shape<TI, BATCH_SIZE*MODULE_SPEC::N_AGENTS, INPUT_DIM/MODULE_SPEC::N_AGENTS>{});
+        auto d_output_inner = reshape_row_major(device, d_output_buffer_view, tensor::Shape<TI, BATCH_SIZE*MODULE_SPEC::N_AGENTS, OUTPUT_DIM/MODULE_SPEC::N_AGENTS>{});
+        backward_full(device, model.content, input_inner, d_output_inner, d_input_inner, buffers.buffer, mode);
+        copy(device, device, d_input_buffer_view, d_input);
+    }
+    template<typename DEVICE, typename MODULE_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, nn_models::multi_agent_wrapper::ModuleBackward<MODULE_SPEC>& model, Tensor<D_OUTPUT_SPEC>& d_output, Tensor<D_INPUT_SPEC>& d_input, nn_models::multi_agent_wrapper::ModuleBuffer<BUFFER_SPEC> buffers, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        using TI = typename DEVICE::index_t;
+        static_assert(D_INPUT_SPEC::SHAPE::LENGTH == 3);
+        static_assert(D_OUTPUT_SPEC::SHAPE::LENGTH == 3);
+        constexpr TI INPUT_DIM = D_INPUT_SPEC::SHAPE::template GET<D_INPUT_SPEC::SHAPE::LENGTH-1>;
+        constexpr TI OUTPUT_DIM = D_OUTPUT_SPEC::SHAPE::template GET<D_OUTPUT_SPEC::SHAPE::LENGTH-1>;
+        static_assert(INPUT_DIM % MODULE_SPEC::N_AGENTS == 0);
+        static_assert(OUTPUT_DIM % MODULE_SPEC::N_AGENTS == 0);
+        constexpr TI BATCH_AXIS = 1;
+        constexpr TI BATCH_SIZE = D_INPUT_SPEC::SHAPE::template GET<BATCH_AXIS>;
+        static_assert(BATCH_SIZE == D_OUTPUT_SPEC::SHAPE::template GET<BATCH_AXIS>);
+        using TI = typename DEVICE::index_t;
+
+        auto d_input_buffer_view = view_range(device, buffers.d_input, tensor::ViewSpec<BATCH_AXIS, BATCH_SIZE>{});
+        auto d_output_buffer_view = view_range(device, buffers.output, tensor::ViewSpec<BATCH_AXIS, BATCH_SIZE>{});
+        copy(device, device, d_output, d_output_buffer_view);
+        auto d_input_inner = reshape_row_major(device, d_input_buffer_view, tensor::Shape<TI, BATCH_SIZE*MODULE_SPEC::N_AGENTS, INPUT_DIM/MODULE_SPEC::N_AGENTS>{});
+        auto d_output_inner = reshape_row_major(device, d_output_buffer_view, tensor::Shape<TI, BATCH_SIZE*MODULE_SPEC::N_AGENTS, OUTPUT_DIM/MODULE_SPEC::N_AGENTS>{});
+        backward_input(device, model.content, d_output_inner, d_input_inner, buffers.buffer, mode);
+        copy(device, device, d_input_buffer_view, d_input);
+    }
+    template<typename DEVICE, typename MODULE_SPEC, typename INPUT_SPEC, typename D_OUTPUT_SPEC, typename BUFFER_SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void backward(DEVICE& device, nn_models::multi_agent_wrapper::ModuleGradient<MODULE_SPEC>& model, const Tensor<INPUT_SPEC>& input, Tensor<D_OUTPUT_SPEC>& d_output, nn_models::multi_agent_wrapper::ModuleBuffer<BUFFER_SPEC> buffers, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        using TI = typename DEVICE::index_t;
+        static_assert(INPUT_SPEC::SHAPE::LENGTH == 3);
+        static_assert(D_OUTPUT_SPEC::SHAPE::LENGTH == 3);
+        constexpr TI INPUT_DIM = INPUT_SPEC::SHAPE::template GET<INPUT_SPEC::SHAPE::LENGTH-1>;
+        constexpr TI OUTPUT_DIM = D_OUTPUT_SPEC::SHAPE::template GET<D_OUTPUT_SPEC::SHAPE::LENGTH-1>;
+        static_assert(INPUT_DIM % MODULE_SPEC::N_AGENTS == 0);
+        static_assert(OUTPUT_DIM % MODULE_SPEC::N_AGENTS == 0);
+        constexpr TI BATCH_AXIS = 1;
+        constexpr TI BATCH_SIZE = INPUT_SPEC::SHAPE::template GET<BATCH_AXIS>;
+        static_assert(BATCH_SIZE == D_OUTPUT_SPEC::SHAPE::template GET<BATCH_AXIS>);
+        using TI = typename DEVICE::index_t;
+
+        auto input_buffer_view = view_range(device, buffers.input, tensor::ViewSpec<BATCH_AXIS, BATCH_SIZE>{});
+        auto d_output_buffer_view = view_range(device, buffers.output, tensor::ViewSpec<BATCH_AXIS, BATCH_SIZE>{});
+        copy(device, device, input, input_buffer_view);
+        copy(device, device, d_output, d_output_buffer_view);
+        auto input_inner = reshape_row_major(device, input_buffer_view, tensor::Shape<TI, BATCH_SIZE*MODULE_SPEC::N_AGENTS, INPUT_DIM/MODULE_SPEC::N_AGENTS>{});
+        auto d_output_inner = reshape_row_major(device, d_output_buffer_view, tensor::Shape<TI, BATCH_SIZE*MODULE_SPEC::N_AGENTS, OUTPUT_DIM/MODULE_SPEC::N_AGENTS>{});
+        backward(device, model.content, input_inner, d_output_inner, buffers.buffer, mode);
+    }
+    template<typename DEVICE, typename SPEC, typename OPTIMIZER>
+    RL_TOOLS_FUNCTION_PLACEMENT void update(DEVICE& device, nn_models::multi_agent_wrapper::ModuleGradient<SPEC>& model, OPTIMIZER& optimizer) {
+        update(device, model.content, 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_models::multi_agent_wrapper::ModuleForward<SOURCE_SPEC>& source, nn_models::multi_agent_wrapper::ModuleForward<TARGET_SPEC>& target){
+        copy(source_device, target_device, source.content, target.content);
+    }
+
+    template<typename DEVICE, typename SPEC_A, typename SPEC_B>
+    RL_TOOLS_FUNCTION_PLACEMENT typename SPEC_A::T abs_diff(DEVICE& device, nn_models::multi_agent_wrapper::ModuleForward<SPEC_A>& a, const nn_models::multi_agent_wrapper::ModuleForward<SPEC_B>& b){
+        auto diff = abs_diff(device, a.content, b.content);
+        return diff;
+    }
+
+
+    template<typename DEVICE, typename MODULE_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT bool is_nan(DEVICE& device, nn_models::multi_agent_wrapper::ModuleForward<MODULE_SPEC>& model){
+        bool current_module_nan = is_nan(device, model.content);
+        return current_module_nan;
+    }
+    template<typename DEVICE, typename SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT bool is_nan(DEVICE& device, nn_models::multi_agent_wrapper::ModuleState<SPEC>& state, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        return is_nan(device, state.inner_state, mode);
+    }
+
+    template <typename DEVICE, typename BUFFER_SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT void sample(DEVICE& device, nn_models::multi_agent_wrapper::ModuleBuffer<BUFFER_SPEC>& buffer, RNG& rng){
+        sample(device, buffer.buffer, rng);
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void print(DEVICE& device, const nn_models::multi_agent_wrapper::ModuleForward<SPEC>& model, typename DEVICE::index_t layer_i = 0){
+        print(device, model.content, layer_i);
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/nn_models/multi_agent_wrapper/persist.h

@@ -0,0 +1,25 @@
+#include "../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_MODELS_MULTI_AGENT_WRAPPER_PERSIST_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_MODELS_MULTI_AGENT_WRAPPER_PERSIST_H
+#include "../../nn/parameters/persist.h"
+#include "../../nn/persist.h"
+#include "model.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template<typename DEVICE, typename SPEC, typename GROUP>
+    void save(DEVICE& device, nn_models::multi_agent_wrapper::ModuleForward<SPEC>& model, GROUP& group) {
+        auto wrapper_group = create_group(device, group, "multi_agent_wrapper");
+        auto content_group = create_group(device, wrapper_group, "content");
+        save(device, model.content, content_group);
+    }
+    template<typename DEVICE, typename SPEC, typename GROUP>
+    bool load(DEVICE& device, nn_models::multi_agent_wrapper::ModuleForward<SPEC>& model, GROUP& group, typename DEVICE::index_t layer_i = 0) {
+        auto wrapper_group = get_group(device, group, "multi_agent_wrapper");
+        auto content_group = get_group(device, wrapper_group, "content");
+        return load(device, model.content, content_group);
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/nn_models/multi_agent_wrapper/persist_code.h

@@ -0,0 +1,71 @@
+#include "../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_MODELS_MULTI_AGENT_WRAPPER_PERSIST_CODE_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_MODELS_MULTI_AGENT_WRAPPER_PERSIST_CODE_H
+
+
+
+#include <string>
+#include <sstream>
+#include "../../persist/code.h"
+#include "model.h"
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template<typename DEVICE, typename SPEC>
+    persist::Code save_code_split(DEVICE& device, nn_models::multi_agent_wrapper::ModuleForward<SPEC>& module, 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();
+        std::stringstream ss_header;
+        ss_header << "#include <rl_tools/nn_models/multi_agent_wrapper/model.h>\n";
+        std::stringstream ss;
+        ss << ind << "namespace " << name << " {\n";
+        auto content = save_code_split(device, module.content, "content", const_declaration, indent+1);
+        ss_header << content.header;
+        ss << content.body;
+        std::string T_string = containers::persist::get_type_string<typename SPEC::TYPE_POLICY::DEFAULT>();
+        std::string TI_string = containers::persist::get_type_string<TI>();
+        ss << "\n";
+        ss << ind << "    using TYPE_POLICY = " << to_string(typename SPEC::TYPE_POLICY{}) << ";\n";
+        ss << ind << "    using CONFIG = RL_TOOLS""_NAMESPACE_WRAPPER ::rl_tools::nn_models::multi_agent_wrapper::Configuration<";
+        ss << "TYPE_POLICY, ";
+        ss << TI_string << ", ";
+        ss << std::to_string(SPEC::N_AGENTS) << ", ";
+        ss << "content::model_definition::MODULE_CHAIN";
+        ss << ">; \n";
+        ss << ind << "    using CAPABILITY = " << to_string(typename SPEC::CAPABILITY::template CHANGE_PARAMETERS<true, true>{}) << "; \n";
+        ss << ind << "    using INPUT_SHAPE = RL_TOOLS""_NAMESPACE_WRAPPER ::rl_tools::tensor::Shape<" << TI_string << ", " << SPEC::INPUT_SHAPE::template GET<0> << ", " << SPEC::INPUT_SHAPE::template GET<1> << ", " << SPEC::INPUT_SHAPE::template GET<2> << ">;\n";
+        ss << ind << "    using TYPE = RL_TOOLS""_NAMESPACE_WRAPPER ::rl_tools::nn_models::multi_agent_wrapper::Module<CONFIG, CAPABILITY, INPUT_SHAPE>; \n";
+        std::stringstream ss_initializer_list, ss_initializer_list_function;
+        ss_initializer_list << "{content::factory<TYPE::MODEL>}";
+        ss_initializer_list_function << "{content::factory_function<TYPE::MODEL>()}";
+        ss << ind << "    " << (const_declaration ? "constexpr " : "") << "TYPE module = " << ss_initializer_list.str() << ";\n";
+        ss << ind << "    " << "template <typename T_TYPE = TYPE>" << "\n";
+        ss << ind << "    " << (const_declaration ? "constexpr " : "") << "T_TYPE factory = " << ss_initializer_list.str() << ";" << "\n";
+        ss << ind << "    " << "template <typename T_TYPE = TYPE>" << "\n";
+        ss << ind << "    " << (const_declaration ? "constexpr " : "") << "T_TYPE factory_function(){return " << ss_initializer_list_function.str() << ";}\n";
+        ss << ind << "}\n";
+        return {ss_header.str(), ss.str()};
+    }
+    template<typename DEVICE, typename SPEC>
+    std::string save_code(DEVICE& device, nn_models::multi_agent_wrapper::ModuleForward<SPEC>& module, std::string name, bool const_declaration = true, typename DEVICE::index_t indent = 0) {
+        auto code = save_code_split(device, module, name, const_declaration, indent);
+        return code.header + code.body;
+    }
+    template <typename DEVICE, typename SPEC>
+    std::string nn_analytics(DEVICE& device, nn_models::multi_agent_wrapper::ModuleGradient<SPEC>& nn) {
+        std::string data;
+        data += "{";
+        data += "   \"type\": \"multi_agent_wrapper\",";
+        data += "   \"inner\": ";
+        data += nn_analytics(device, nn.content);
+        data += "}";
+        return data;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/nn_models/operations_cuda.h

@@ -0,0 +1 @@
+#include "mlp/operations_cuda.h"

hyperdrone/_vendor/rl-tools/include/rl_tools/nn_models/operations_dummy.h

@@ -0,0 +1 @@
+#include "mlp/operations_dummy.h"

hyperdrone/_vendor/rl-tools/include/rl_tools/nn_models/operations_generic.h

@@ -0,0 +1,4 @@
+#include "../rl_tools.h"
+
+#include "mlp/operations_generic.h"
+