hyperdrone 0.1.0__py3-none-any.whl

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

@@ -0,0 +1,76 @@
+#include "../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_MODELS_MLP_UNCONDITIONAL_STDDEV_OPERATIONS_GENERIC_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_MODELS_MLP_UNCONDITIONAL_STDDEV_OPERATIONS_GENERIC_H
+
+#include "network.h"
+#include "../../nn_models/mlp/operations_generic.h"
+
+
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template <typename DEVICE, typename SPEC, template <typename> typename BASE>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn_models::mlp_unconditional_stddev::NeuralNetworkForward<SPEC, BASE>& m){
+        malloc(device, static_cast<nn_models::mlp::NeuralNetworkForward<SPEC>&>(m));
+        malloc(device, m.log_std);
+    }
+    template <typename DEVICE, typename SPEC, template <typename> typename BASE>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn_models::mlp_unconditional_stddev::NeuralNetworkForward<SPEC, BASE>& m){
+        free(device, static_cast<nn_models::mlp::NeuralNetworkForward<SPEC>&>(m));
+        free(device, m.log_std);
+    }
+    template <typename DEVICE, typename SPEC, template <typename> typename BASE, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT void init_weights(DEVICE& device, nn_models::mlp_unconditional_stddev::NeuralNetworkForward<SPEC, BASE>& m, RNG& rng){
+        init_weights(device, static_cast<nn_models::mlp::NeuralNetworkForward<SPEC>&>(m), rng);
+        set_all(device, m.log_std.parameters, 0);
+    }
+    template<typename DEVICE, typename SPEC, template <typename> typename BASE, typename OPTIMIZER>
+    RL_TOOLS_FUNCTION_PLACEMENT void update(DEVICE& device, nn_models::mlp_unconditional_stddev::NeuralNetworkGradient<SPEC, BASE>& network, OPTIMIZER& optimizer) {
+        update(device, network.log_std, optimizer);
+        update(device, static_cast<nn_models::mlp::NeuralNetworkGradient<SPEC>&>(network), optimizer);
+    }
+    template<typename DEVICE, typename SPEC, template <typename> typename BASE>
+    RL_TOOLS_FUNCTION_PLACEMENT void zero_gradient(DEVICE& device, nn_models::mlp_unconditional_stddev::NeuralNetworkGradient<SPEC, BASE>& network) {
+        zero_gradient(device, static_cast<nn_models::mlp::NeuralNetworkGradient<SPEC>&>(network));
+        zero_gradient(device, network.log_std);
+    }
+
+    template<typename DEVICE, typename SPEC, template <typename> typename BASE, typename OPTIMIZER>
+    RL_TOOLS_FUNCTION_PLACEMENT void _reset_optimizer_state(DEVICE& device, nn_models::mlp_unconditional_stddev::NeuralNetworkGradient<SPEC, BASE>& network, OPTIMIZER& optimizer) {
+        _reset_optimizer_state(device, static_cast<nn_models::mlp::NeuralNetworkGradient<SPEC>&>(network), optimizer);
+        _reset_optimizer_state(device, network.log_std, optimizer);
+    }
+
+    template<typename SOURCE_DEVICE, typename TARGET_DEVICE, typename SOURCE_SPEC, typename TARGET_SPEC, template <typename> typename SOURCE_BASE, template <typename> typename TARGET_BASE>
+    RL_TOOLS_FUNCTION_PLACEMENT void copy(SOURCE_DEVICE& source_device, TARGET_DEVICE& target_device, const  nn_models::mlp_unconditional_stddev::NeuralNetworkForward<SOURCE_SPEC, SOURCE_BASE>& source, nn_models::mlp_unconditional_stddev::NeuralNetworkForward<TARGET_SPEC, TARGET_BASE>& target){
+        static_assert(rl_tools::nn_models::mlp::check_spec_memory<SOURCE_SPEC, TARGET_SPEC>, "The source and target network must have the same structure");
+        copy(source_device, target_device, static_cast<const nn_models::mlp::NeuralNetworkForward<SOURCE_SPEC>&>(source), static_cast<nn_models::mlp::NeuralNetworkForward<TARGET_SPEC>&>(target));
+        copy(source_device, target_device, source.log_std, target.log_std);
+    }
+    template<typename SOURCE_DEVICE, typename TARGET_DEVICE, typename SOURCE_SPEC, typename TARGET_SPEC, typename SOURCE, template <typename> typename TARGET_BASE>
+    RL_TOOLS_FUNCTION_PLACEMENT void copy_from_generic(SOURCE_DEVICE& source_device, TARGET_DEVICE& target_device, const SOURCE& source, nn_models::mlp_unconditional_stddev::NeuralNetworkForward<TARGET_SPEC, TARGET_BASE>& target){
+        static_assert(rl_tools::nn_models::mlp::check_spec_memory<SOURCE_SPEC, TARGET_SPEC>, "The source and target network must have the same structure");
+        copy_from_generic(source_device, target_device, static_cast<const nn_models::mlp::NeuralNetworkForward<SOURCE_SPEC>&>(source), static_cast<nn_models::mlp::NeuralNetworkForward<TARGET_SPEC>&>(target));
+        copy_from_generic(source_device, target_device, source.log_std, target.log_std);
+    }
+    template <typename DEVICE, typename SPEC_1, typename SPEC_2, template <typename> typename BASE_1, template <typename> typename BASE_2>
+    RL_TOOLS_FUNCTION_PLACEMENT typename SPEC_1::TYPE_POLICY::DEFAULT abs_diff(DEVICE& device, nn_models::mlp_unconditional_stddev::NeuralNetworkForward<SPEC_1, BASE_1>& n1, const nn_models::mlp_unconditional_stddev::NeuralNetworkForward<SPEC_2, BASE_2>& n2){
+        using T = typename SPEC_1::TYPE_POLICY::DEFAULT;
+        T acc = 0;
+        acc += abs_diff(device, static_cast<BASE_1<SPEC_1>&>(n1), static_cast<const BASE_2<SPEC_2>&>(n2));
+        acc += abs_diff(device, n1.log_std, n2.log_std);
+        return acc;
+    }
+    template <typename DEVICE, typename SPEC_1, typename SPEC_2, template <typename> typename BASE_1, template <typename> typename BASE_2>
+    RL_TOOLS_FUNCTION_PLACEMENT typename SPEC_1::TYPE_POLICY::DEFAULT abs_diff(DEVICE& device, nn_models::mlp_unconditional_stddev::NeuralNetworkBackward<SPEC_1, BASE_1>& n1, const nn_models::mlp_unconditional_stddev::NeuralNetworkBackward<SPEC_2, BASE_2>& n2){
+        return abs_diff(device, static_cast<nn_models::mlp_unconditional_stddev::NeuralNetworkForward<SPEC_1, BASE_1>&>(n1), static_cast<const nn_models::mlp_unconditional_stddev::NeuralNetworkForward<SPEC_2, BASE_2>&>(n2));
+    }
+    template <typename DEVICE, typename SPEC_1, typename SPEC_2, template <typename> typename BASE_1, template <typename> typename BASE_2>
+    RL_TOOLS_FUNCTION_PLACEMENT typename SPEC_1::TYPE_POLICY::DEFAULT abs_diff(DEVICE& device, nn_models::mlp_unconditional_stddev::NeuralNetworkGradient<SPEC_1, BASE_1>& n1, const nn_models::mlp_unconditional_stddev::NeuralNetworkGradient<SPEC_2, BASE_2>& n2){
+        return abs_diff(device, static_cast<nn_models::mlp_unconditional_stddev::NeuralNetworkBackward<SPEC_1, BASE_1>&>(n1), static_cast<const nn_models::mlp_unconditional_stddev::NeuralNetworkBackward<SPEC_2, BASE_2>&>(n2));
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif

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

@@ -0,0 +1,43 @@
+#include "../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_MODELS_MLP_UNCONDITIONAL_STDDEV_PERSIST_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_MODELS_MLP_UNCONDITIONAL_STDDEV_PERSIST_H
+#include "../../nn/parameters/persist.h"
+#include "../../nn/persist.h"
+#include "network.h"
+#include "../mlp/persist.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template<typename DEVICE, typename SPEC, template <typename> typename BASE, typename GROUP>
+    void save(DEVICE& device, nn_models::mlp_unconditional_stddev::NeuralNetworkForward<SPEC, BASE>& network, GROUP& group) {
+        save(device, static_cast<nn_models::mlp::NeuralNetworkForward<SPEC>&>(network), group);
+        auto log_std_group = create_group(device, group, "log_std");
+        save(device, network.log_std, log_std_group);
+    }
+    template<typename DEVICE, typename SPEC, template <typename> typename BASE, typename GROUP>
+    void save(DEVICE& device, nn_models::mlp_unconditional_stddev::NeuralNetworkBackward<SPEC, BASE>& network, GROUP& group){
+        save(device, static_cast<nn_models::mlp_unconditional_stddev::NeuralNetworkForward<SPEC, BASE>&>(network), group);
+    }
+    template<typename DEVICE, typename SPEC, template <typename> typename BASE, typename GROUP>
+    void save(DEVICE& device, nn_models::mlp_unconditional_stddev::NeuralNetworkGradient<SPEC, BASE>& network, GROUP& group){
+        save(device, static_cast<nn_models::mlp_unconditional_stddev::NeuralNetworkBackward<SPEC, BASE>&>(network), group);
+    }
+    template<typename DEVICE, typename SPEC, template <typename> typename BASE, typename GROUP>
+    bool load(DEVICE& device, nn_models::mlp_unconditional_stddev::NeuralNetworkForward<SPEC, BASE>& network, GROUP& group){
+        bool success = load(device, static_cast<nn_models::mlp::NeuralNetworkForward<SPEC>&>(network), group);
+        auto log_std_group = get_group(device, group, "log_std");
+        success &= load(device, network.log_std, log_std_group);
+        return success;
+    }
+    template<typename DEVICE, typename SPEC, template <typename> typename BASE, typename GROUP>
+    bool load(DEVICE& device, nn_models::mlp_unconditional_stddev::NeuralNetworkBackward<SPEC, BASE>& network, GROUP& group){
+        return load(device, static_cast<nn_models::mlp_unconditional_stddev::NeuralNetworkForward<SPEC, BASE>&>(network), group);
+    }
+    template<typename DEVICE, typename SPEC, template <typename> typename BASE, typename GROUP>
+    bool load(DEVICE& device, nn_models::mlp_unconditional_stddev::NeuralNetworkGradient<SPEC, BASE>& network, GROUP& group){
+        return load(device, static_cast<nn_models::mlp_unconditional_stddev::NeuralNetworkBackward<SPEC, BASE>&>(network), group);
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

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

@@ -0,0 +1,135 @@
+#include "../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_MODELS_MLP_UNCONDITIONAL_STDDEV_PERSIST_CODE_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_MODELS_MLP_UNCONDITIONAL_STDDEV_PERSIST_CODE_H
+
+
+
+#include <string>
+#include <sstream>
+#include "../../persist/code.h"
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template<typename DEVICE, typename SPEC, template <typename> typename BASE>
+    persist::Code save_code_split(DEVICE& device, nn_models::mlp_unconditional_stddev::NeuralNetworkForward<SPEC, BASE>& network, std::string name, bool const_declaration=true, typename DEVICE::index_t indent = 0){
+        // using T = typename SPEC::T;
+        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;
+        std::stringstream ss;
+        ss << ind << "namespace " << name << " {\n";
+        auto input_layer = save_code_split(device, network.input_layer, "input_layer", const_declaration, indent+1);
+        ss_header << input_layer.header;
+        ss_header << "#include <rl_tools/nn_models/mlp_unconditional_stddev/network.h>\n";
+        ss << input_layer.body;
+        for(TI hidden_layer_i = 0; hidden_layer_i < SPEC::NUM_HIDDEN_LAYERS; hidden_layer_i++){
+            auto hidden_layer = save_code_split(device, network.hidden_layers[hidden_layer_i], "hidden_layer_" + std::to_string(hidden_layer_i), const_declaration, indent+1);
+            ss_header << hidden_layer.header;
+            ss << hidden_layer.body;
+        }
+        auto output_layer = save_code_split(device, network.output_layer, "output_layer", const_declaration, indent+1);
+        ss_header << output_layer.header;
+        ss << output_layer.body;
+        auto log_std = save_code_split(device, network.log_std, "log_std", const_declaration, indent+1);
+        ss_header << log_std.header;
+        ss << log_std.body;
+        // std::string T_string = containers::persist::get_type_string<T>();
+        std::string TI_string = containers::persist::get_type_string<TI>();
+        ss << ind << "    using TYPE_POLICY = " + to_string(typename SPEC::TYPE_POLICY{}) + ";" << "\n";
+        ss << ind << "    using CONFIG = RL_TOOLS""_NAMESPACE_WRAPPER ::rl_tools::nn_models::mlp::Configuration<";
+        ss << "TYPE_POLICY, ";
+        ss << TI_string << ", ";
+        ss << SPEC::OUTPUT_DIM << ", " << SPEC::NUM_LAYERS << ", " << SPEC::HIDDEN_DIM << ", ";
+        ss << nn::layers::dense::persist::get_activation_function_string<SPEC::HIDDEN_ACTIVATION_FUNCTION>() << ", ";
+        ss << nn::layers::dense::persist::get_activation_function_string<SPEC::OUTPUT_ACTIVATION_FUNCTION>() << ", ";
+        ss << "RL_TOOLS""_NAMESPACE_WRAPPER ::rl_tools::nn::layers::dense::DefaultInitializer<" << "TYPE_POLICY, " << TI_string << ">";
+        ss << ">; \n";
+        ss << ind << "    " << "using TEMPLATE = RL_TOOLS""_NAMESPACE_WRAPPER ::rl_tools::nn_models::mlp_unconditional_stddev::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::template CHANGE_PARAMETERS<true, true>{}) << ";" << "\n";
+        ss << ind << "    " << "using TYPE = RL_TOOLS""_NAMESPACE_WRAPPER ::rl_tools::nn_models::mlp_unconditional_stddev::NeuralNetwork<CONFIG, CAPABILITY, INPUT_SHAPE>;" << "\n";
+        std::stringstream ss_initializer_list;
+        {
+            ss_initializer_list << "{{input_layer::factory<TYPE::SPEC::INPUT_LAYER>, ";
+            ss_initializer_list << "{";
+            for(TI hidden_layer_i = 0; hidden_layer_i < SPEC::NUM_HIDDEN_LAYERS; hidden_layer_i++){
+                if(hidden_layer_i > 0){
+                    ss_initializer_list << ", ";
+                }
+                ss_initializer_list << "hidden_layer_" << hidden_layer_i << "::factory<TYPE::SPEC::HIDDEN_LAYER>";
+            }
+            ss_initializer_list << "}, ";
+            ss_initializer_list << "output_layer::factory<TYPE::SPEC::OUTPUT_LAYER>}";
+            ss_initializer_list << ", log_std::parameters}";
+        }
+        std::stringstream ss_initializer_list_create, ss_initializer_list_create_function;
+        {
+            ss_initializer_list_create << "{{input_layer::factory<typename T_TYPE::SPEC::INPUT_LAYER>, ";
+            ss_initializer_list_create_function << "{{input_layer::factory_function<typename T_TYPE::SPEC::INPUT_LAYER>(), ";
+            ss_initializer_list_create << "{";
+            ss_initializer_list_create_function << "{";
+            for(TI hidden_layer_i = 0; hidden_layer_i < SPEC::NUM_HIDDEN_LAYERS; hidden_layer_i++){
+                if(hidden_layer_i > 0){
+                    ss_initializer_list_create << ", ";
+                    ss_initializer_list_create_function << ", ";
+                }
+                ss_initializer_list_create << "hidden_layer_" << hidden_layer_i << "::factory<typename T_TYPE::SPEC::HIDDEN_LAYER>";
+                ss_initializer_list_create_function << "hidden_layer_" << hidden_layer_i << "::factory_function<typename T_TYPE::SPEC::HIDDEN_LAYER>()";
+            }
+            ss_initializer_list_create << "}, ";
+            ss_initializer_list_create_function << "}, ";
+            ss_initializer_list_create << "output_layer::factory<typename T_TYPE::SPEC::OUTPUT_LAYER>}";
+            ss_initializer_list_create_function << "output_layer::factory_function<typename T_TYPE::SPEC::OUTPUT_LAYER>()";
+            ss_initializer_list_create << ", log_std::parameters}";
+            ss_initializer_list_create_function << ", log_std::parameters}";
+        }
+
+        std::string initializer_list, initializer_list_create;
+        if constexpr(SPEC::CAPABILITY::TAG == nn::LayerCapability::Forward){
+            initializer_list = ss_initializer_list.str();
+            initializer_list_create = ss_initializer_list_create.str();
+        }
+        else{
+            if constexpr(SPEC::CAPABILITY::TAG == nn::LayerCapability::Backward){
+                initializer_list = "{" + ss_initializer_list.str() + "}";
+                initializer_list_create = "{" + ss_initializer_list_create.str() + "}";
+            }
+            else{
+                if constexpr(SPEC::CAPABILITY::TAG == nn::LayerCapability::Gradient){
+                    initializer_list = "{{" + ss_initializer_list.str() + "}}";
+                    initializer_list_create = "{{" + ss_initializer_list_create.str() + "}}";
+                }
+                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_create << ";" << "\n";
+        ss << ind << "    " << "template <typename T_TYPE = TYPE>" << "\n";
+        ss << ind << "    " << (const_declaration ? "constexpr " : "") << "T_TYPE factory_function(){return T_TYPE" << initializer_list_create << ";" << "}\n";
+        ss << ind << "}\n";
+        return {ss_header.str(), ss.str()};
+    }
+    template<typename DEVICE, typename SPEC, template <typename> typename BASE>
+    persist::Code save_code_split(DEVICE& device, nn_models::mlp_unconditional_stddev::NeuralNetworkBackward<SPEC, BASE>& network, std::string name, bool const_declaration=true, typename DEVICE::index_t indent = 0){
+        return save_code_split(device, static_cast<nn_models::mlp_unconditional_stddev::NeuralNetworkForward<SPEC, BASE>&>(network), name, const_declaration, indent);
+    }
+    template<typename DEVICE, typename SPEC, template <typename> typename BASE>
+    persist::Code save_code_split(DEVICE& device, nn_models::mlp_unconditional_stddev::NeuralNetworkGradient<SPEC, BASE>& network, std::string name, bool const_declaration=true, typename DEVICE::index_t indent = 0){
+        return save_code_split(device, static_cast<nn_models::mlp_unconditional_stddev::NeuralNetworkBackward<SPEC, BASE>&>(network), name, const_declaration, indent);
+    }
+    template<typename DEVICE, typename SPEC>
+    std::string save_code(DEVICE& device, nn_models::mlp_unconditional_stddev::NeuralNetworkForward<SPEC>& network, std::string name, bool const_declaration = true, typename DEVICE::index_t indent = 0) {
+        auto code = save_code_split(device, network, name, const_declaration, indent);
+        return code.header + code.body;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif

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

@@ -0,0 +1,3 @@
+#include "mlp/network.h"
+#include "mlp_unconditional_stddev/network.h"
+#include "sequential/model.h"

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