hyperray 0.0.1__py3-none-any.whl

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/algorithms/ppo/operations_generic.h

@@ -0,0 +1,332 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ALGORITHMS_PPO_OPERATIONS_GENERIC_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ALGORITHMS_PPO_OPERATIONS_GENERIC_H
+
+#include "ppo.h"
+#include "../../../nn/loss_functions/mse/operations_generic.h"
+#include "../../../rl/components/on_policy_runner/on_policy_runner.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, rl::algorithms::ppo::Buffers<SPEC>& buffers){
+        malloc(device, buffers.current_batch_actions);
+        malloc(device, buffers.d_critic_output);
+        malloc(device, buffers.d_action_log_prob_d_action);
+        malloc(device, buffers.d_action_log_prob_d_action_log_std);
+        malloc(device, buffers.rollout_log_std);
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, rl::algorithms::ppo::Buffers<SPEC>& buffers){
+        free(device, buffers.current_batch_actions);
+        free(device, buffers.d_critic_output);
+        free(device, buffers.d_action_log_prob_d_action);
+        free(device, buffers.d_action_log_prob_d_action_log_std);
+        free(device, buffers.rollout_log_std);
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, rl::algorithms::PPO<SPEC>& ppo){
+        malloc(device, ppo.actor);
+        malloc(device, ppo.critic);
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, rl::algorithms::PPO<SPEC>& ppo){
+        free(device, ppo.actor);
+        free(device, ppo.critic);
+    }
+    template <typename DEVICE, typename SPEC, typename ACTOR_OPTIMIZER, typename CRITIC_OPTIMIZER, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT void init(DEVICE& device, rl::algorithms::PPO<SPEC>& ppo, ACTOR_OPTIMIZER& actor_optimizer, CRITIC_OPTIMIZER& critic_optimizer, RNG& rng){
+#ifdef RL_TOOLS_DEBUG_RL_ALGORITHMS_PPO_CHECK_INIT
+        ppo.initialized = true;
+#endif
+        init_weights(device, ppo.actor, rng);
+        init(device, actor_optimizer);
+        reset_forward_state(device, ppo.actor);
+        zero_gradient(device, ppo.actor);
+        reset_optimizer_state(device, actor_optimizer, ppo.actor);
+        auto& last_layer = get_last_layer(ppo.actor);
+        set_all(device, last_layer.log_std.parameters, math::log(device.math, SPEC::PARAMETERS::INITIAL_ACTION_STD));
+        init_weights(device, ppo.critic, rng);
+        init(device, critic_optimizer);
+        reset_forward_state(device, ppo.actor);
+        zero_gradient(device, ppo.actor);
+        reset_optimizer_state(device, critic_optimizer, ppo.critic);
+    }
+    template <typename DEVICE, typename DATASET_SPEC, typename BOOTSTRAP_SPEC, typename PPO_PARAMETERS>
+    RL_TOOLS_FUNCTION_PLACEMENT void estimate_generalized_advantages(DEVICE& device, rl::components::on_policy_runner::Dataset<DATASET_SPEC>& dataset, const Matrix<BOOTSTRAP_SPEC>& bootstrap_values, PPO_PARAMETERS ppo_parameters_tag){
+        static_assert(DATASET_SPEC::SPEC::COLLECT_NEXT_OBSERVATIONS || !(PPO_PARAMETERS::BOOTSTRAP_TRUNCATIONS || PPO_PARAMETERS::IGNORE_TERMINATION), "PPO bootstrapping requires pre-reset observation capture");
+        static_assert(BOOTSTRAP_SPEC::ROWS == DATASET_SPEC::STEPS_TOTAL && BOOTSTRAP_SPEC::COLS == 1, "GAE requires one bootstrap value per transition");
+        using OPR_SPEC = typename DATASET_SPEC::SPEC;
+        using BUFFER = decltype(dataset);
+        using T = typename DATASET_SPEC::SPEC::TYPE_POLICY::DEFAULT;
+        using TI = typename DEVICE::index_t;
+        constexpr TI STEPS_PER_ENV = DATASET_SPEC::STEPS_PER_ENV;
+        for(TI env_i = 0; env_i < OPR_SPEC::N_ENVIRONMENTS; env_i++){
+            T previous_advantage = 0;
+            for(TI step_forward_i = 0; step_forward_i < STEPS_PER_ENV; step_forward_i++){
+                TI step_backward_i = (STEPS_PER_ENV - 1 - step_forward_i);
+                TI pos = step_backward_i * OPR_SPEC::N_ENVIRONMENTS + env_i;
+                bool terminated = get(dataset.terminated, pos, 0);
+                bool truncated = get(dataset.truncated, pos, 0);
+#ifdef RL_TOOLS_DEBUG_RL_ALGORITHMS_PPO_GAE_CHECK_TERMINATED_TRUNCATED
+                utils::assert_exit(device, !terminated || (terminated && truncated), "terminationn should imply truncation");
+#endif
+                T current_step_value = get(dataset.values, pos, 0);
+                bool stop_bootstrap = terminated ? !PPO_PARAMETERS::IGNORE_TERMINATION : (truncated && !PPO_PARAMETERS::BOOTSTRAP_TRUNCATIONS);
+                T next_step_value = stop_bootstrap ? 0 : get(bootstrap_values, pos, 0);
+
+                T td_error = get(dataset.rewards, pos, 0) + PPO_PARAMETERS::GAMMA * next_step_value - current_step_value;
+                if(truncated){
+                    previous_advantage = 0;
+                }
+                T advantage = PPO_PARAMETERS::LAMBDA * PPO_PARAMETERS::GAMMA * previous_advantage + td_error;
+                set(dataset.advantages, pos, 0, advantage);
+                set(dataset.target_values, pos, 0, advantage + current_step_value);
+                previous_advantage = advantage;
+            }
+        }
+    }
+    template <typename DEVICE, typename PPO_SPEC, typename BUFFERS_SPEC, typename BATCH_ACTIONS_SPEC, typename BATCH_ACTIONS_MEAN_SPEC, typename BATCH_ACTION_LOG_PROBS_SPEC, typename BATCH_ADVANTAGES_SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT void ppo_compute_actor_loss_gradient(DEVICE& device, rl::algorithms::PPO<PPO_SPEC>& ppo, rl::algorithms::ppo::Buffers<BUFFERS_SPEC>& ppo_buffers, Matrix<BATCH_ACTIONS_SPEC>& batch_actions, Matrix<BATCH_ACTIONS_MEAN_SPEC>& batch_actions_mean, Matrix<BATCH_ACTION_LOG_PROBS_SPEC>& batch_action_log_probs, Matrix<BATCH_ADVANTAGES_SPEC>& batch_advantages, typename PPO_SPEC::TYPE_POLICY::DEFAULT& policy_kl_divergence, typename PPO_SPEC::TYPE_POLICY::DEFAULT& batch_policy_kl_divergence, RNG& rng){
+        using T = typename PPO_SPEC::TYPE_POLICY::DEFAULT;
+        using TI = typename PPO_SPEC::TI;
+        T advantage_mean = 0;
+        T advantage_std = 0;
+        if(PPO_SPEC::PARAMETERS::NORMALIZE_ADVANTAGE){
+            for(TI batch_step_i = 0; batch_step_i < PPO_SPEC::PARAMETERS::BATCH_SIZE; batch_step_i++){
+                T advantage = get(batch_advantages, batch_step_i, 0);
+                advantage_mean += advantage;
+                advantage_std += advantage * advantage;
+            }
+            advantage_mean /= PPO_SPEC::PARAMETERS::BATCH_SIZE;
+            advantage_std /= PPO_SPEC::PARAMETERS::BATCH_SIZE;
+            advantage_std = math::sqrt(device.math, math::max(device.math, (T)0, advantage_std - advantage_mean * advantage_mean));
+        }
+        constexpr TI BATCH_SIZE = PPO_SPEC::PARAMETERS::BATCH_SIZE;
+        constexpr TI ACTION_DIM = PPO_SPEC::ENVIRONMENT::ACTION_DIM;
+        constexpr TI N_AGENTS = PPO_SPEC::ENVIRONMENT::N_AGENTS;
+        constexpr TI PER_AGENT_ACTION_DIM = ACTION_DIM / N_AGENTS;
+        for(TI batch_step_i = 0; batch_step_i < BATCH_SIZE; batch_step_i++){
+            T action_log_prob = 0;
+            for(TI action_i = 0; action_i < ACTION_DIM; action_i++){
+                T current_action = get(ppo_buffers.current_batch_actions, batch_step_i, action_i);
+                T rollout_action = get(batch_actions, batch_step_i, action_i);
+                auto& last_layer = get_last_layer(ppo.actor);
+                T current_action_log_std = get(device, last_layer.log_std.parameters, action_i % PER_AGENT_ACTION_DIM);
+                T current_action_std = math::exp(device.math, current_action_log_std);
+                if(PPO_SPEC::PARAMETERS::ADAPTIVE_LEARNING_RATE){
+                    T rollout_action_log_std = get(ppo_buffers.rollout_log_std, 0, action_i);
+                    T rollout_action_std = math::exp(device.math, rollout_action_log_std);
+                    T rollout_action_mean = get(batch_actions_mean, batch_step_i, action_i);
+                    T action_mean_diff = rollout_action_mean - current_action;
+                    T kl = rollout_action_log_std - current_action_log_std;
+                    kl += (current_action_std * current_action_std + action_mean_diff * action_mean_diff)/(2 * rollout_action_std * rollout_action_std + PPO_SPEC::PARAMETERS::POLICY_KL_EPSILON);
+                    kl += (T)-0.5;
+                    kl = math::max(device.math, kl, (T)0);
+                    policy_kl_divergence += kl;
+                    batch_policy_kl_divergence += kl;
+                }
+                action_log_prob += random::normal_distribution::log_prob(device.random, current_action, current_action_log_std, rollout_action);
+                set(ppo_buffers.d_action_log_prob_d_action, batch_step_i, action_i, random::normal_distribution::d_log_prob_d_mean(device.random, current_action, current_action_log_std, rollout_action));
+                T current_entropy = current_action_log_std + math::log(device.math, 2 * math::PI<T>)/(T)2 + (T)1/(T)2;
+                T current_entropy_loss = -(T)1/BATCH_SIZE * PPO_SPEC::PARAMETERS::ACTION_ENTROPY_COEFFICIENT * current_entropy;
+                if(PPO_SPEC::PARAMETERS::LEARN_ACTION_STD){
+                    T d_entropy_loss_d_current_action_log_std = -(T)1/BATCH_SIZE * PPO_SPEC::PARAMETERS::ACTION_ENTROPY_COEFFICIENT;
+                    auto& last_layer2 = get_last_layer(ppo.actor);
+                    increment(device, last_layer2.log_std.gradient, d_entropy_loss_d_current_action_log_std, action_i % PER_AGENT_ACTION_DIM);
+                    T d_action_log_prob_d_current_action_log_std = random::normal_distribution::d_log_prob_d_log_std(device.random, current_action, current_action_log_std, rollout_action);
+                    set(ppo_buffers.d_action_log_prob_d_action_log_std, batch_step_i, action_i, d_action_log_prob_d_current_action_log_std);
+                }
+            }
+            T rollout_action_log_prob = get(batch_action_log_probs, batch_step_i, 0);
+            T advantage = get(batch_advantages, batch_step_i, 0);
+            if(PPO_SPEC::PARAMETERS::NORMALIZE_ADVANTAGE){
+                advantage = (advantage - advantage_mean) / (advantage_std + PPO_SPEC::PARAMETERS::ADVANTAGE_EPSILON);
+            }
+            T log_ratio = action_log_prob - rollout_action_log_prob;
+            T ratio = math::exp(device.math, log_ratio);
+            T clipped_ratio = math::clamp(device.math, ratio, 1 - PPO_SPEC::PARAMETERS::EPSILON_CLIP, 1 + PPO_SPEC::PARAMETERS::EPSILON_CLIP);
+            bool clipped = ratio != clipped_ratio;
+            T normal_advantage = ratio * advantage;
+            T clipped_advantage = clipped_ratio * advantage;
+            T slippage = 0.0;
+            bool ratio_min_switch = normal_advantage - clipped_advantage <= slippage;
+            T d_loss_d_pessimistic_surrogate = -(T)1/BATCH_SIZE;
+            T d_pessimistic_surrogate_d_normal_advantage = ratio_min_switch ? 1 : 0;
+            T d_pessimistic_surrogate_d_clipped_advantage = ratio_min_switch ? 0 : 1;
+            T d_normal_advantage_d_ratio = advantage;
+            T d_clipped_advantage_d_clipped_ratio = advantage;
+            T d_clipped_ratio_d_ratio = clipped ? 0 : 1;
+            T d_pessimistic_surrogate_d_ratio = d_pessimistic_surrogate_d_normal_advantage * d_normal_advantage_d_ratio + d_pessimistic_surrogate_d_clipped_advantage * d_clipped_advantage_d_clipped_ratio * d_clipped_ratio_d_ratio;
+            T d_loss_d_ratio = d_loss_d_pessimistic_surrogate * d_pessimistic_surrogate_d_ratio;
+            T d_ratio_d_action_log_prob = ratio;
+            T d_loss_d_action_log_prob = d_loss_d_ratio * d_ratio_d_action_log_prob;
+            for(TI action_i = 0; action_i < ACTION_DIM; action_i++){
+                multiply(ppo_buffers.d_action_log_prob_d_action, batch_step_i, action_i, d_loss_d_action_log_prob);
+                if(PPO_SPEC::PARAMETERS::LEARN_ACTION_STD){
+                    T current_d_action_log_prob_d_action_log_std = get(ppo_buffers.d_action_log_prob_d_action_log_std, batch_step_i, action_i);
+                    auto& last_layer3 = get_last_layer(ppo.actor);
+                    increment(device, last_layer3.log_std.gradient, d_loss_d_action_log_prob * current_d_action_log_prob_d_action_log_std, action_i % PER_AGENT_ACTION_DIM);
+                }
+            }
+        }
+    }
+
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template <typename DEVICE, typename PPO_SPEC, typename DATASET_SPEC, typename ACTOR_OPTIMIZER, typename CRITIC_OPTIMIZER, typename BUFFERS_SPEC, typename ACTOR_BUFFER, typename CRITIC_BUFFER, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT void train(DEVICE& device, rl::algorithms::PPO<PPO_SPEC>& ppo, rl::components::on_policy_runner::Dataset<DATASET_SPEC>& dataset, ACTOR_OPTIMIZER& actor_optimizer, CRITIC_OPTIMIZER& critic_optimizer, rl::algorithms::ppo::Buffers<BUFFERS_SPEC>& ppo_buffers, ACTOR_BUFFER& actor_buffers, CRITIC_BUFFER& critic_buffers, RNG& rng){
+#ifdef RL_TOOLS_DEBUG_RL_ALGORITHMS_PPO_CHECK_INIT
+        utils::assert_exit(device, ppo.initialized, "PPO not initialized");
+#endif
+        using T = typename PPO_SPEC::TYPE_POLICY::DEFAULT;
+        using TI = typename PPO_SPEC::TI;
+        static_assert(PPO_SPEC::ENVIRONMENT::ACTION_DIM == DATASET_SPEC::SPEC::BATCH_ENVIRONMENT::ACTION_DIM, "environment action dimension mismatch");
+        static_assert(PPO_SPEC::ENVIRONMENT::N_AGENTS == DATASET_SPEC::SPEC::BATCH_ENVIRONMENT::N_AGENTS, "environment agent count mismatch");
+        using ENVIRONMENT = typename PPO_SPEC::ENVIRONMENT;
+        using DATASET = rl::components::on_policy_runner::Dataset<DATASET_SPEC>;
+        static_assert(DATASET::STEPS_TOTAL > 1);
+        constexpr TI N_EPOCHS = PPO_SPEC::PARAMETERS::N_EPOCHS;
+        constexpr TI BATCH_SIZE = PPO_SPEC::PARAMETERS::BATCH_SIZE;
+        constexpr TI N_BATCHES = DATASET::STEPS_TOTAL/BATCH_SIZE;
+        static_assert(N_BATCHES > 0);
+        constexpr TI ACTION_DIM = ENVIRONMENT::ACTION_DIM;
+        constexpr TI N_AGENTS = PPO_SPEC::ENVIRONMENT::N_AGENTS;
+        static_assert(ACTION_DIM % N_AGENTS == 0);
+        constexpr TI PER_AGENT_ACTION_DIM = PPO_SPEC::ENVIRONMENT::ACTION_DIM/N_AGENTS;
+
+        using OBS_SHAPE = typename DATASET::OBS_SHAPE;
+        using OBS_PRIV_SHAPE = typename DATASET::OBS_PRIV_SHAPE;
+
+        T policy_kl_divergence = 0;
+        if(PPO_SPEC::PARAMETERS::ADAPTIVE_LEARNING_RATE) {
+            auto& last_layer = get_last_layer(ppo.actor);
+            auto log_std = matrix_view(device, last_layer.log_std.parameters);
+            copy(device, device, log_std, ppo_buffers.rollout_log_std);
+        }
+        for(TI epoch_i = 0; epoch_i < N_EPOCHS; epoch_i++){
+            static_assert(!PPO_SPEC::PARAMETERS::STATEFUL_ACTOR_AND_CRITIC || (N_EPOCHS == 1), "Stateful actor and critic implies single epoch");
+            static_assert(!PPO_SPEC::PARAMETERS::STATEFUL_ACTOR_AND_CRITIC || (PPO_SPEC::PARAMETERS::TRUNCATE_ON_EACH_ITERATION == true), "Stateful actor and critic implies that the OnPolicyRunner should truncate in the beginning of each iteration, to prevent hidden state spillover.");
+            static_assert(!PPO_SPEC::PARAMETERS::STATEFUL_ACTOR_AND_CRITIC || (BATCH_SIZE == DATASET_SPEC::STEPS_PER_ENV * DATASET_SPEC::SPEC::N_ENVIRONMENTS), "Stateful actor and critic implies single batch");
+            static_assert(!PPO_SPEC::PARAMETERS::STATEFUL_ACTOR_AND_CRITIC || !PPO_SPEC::PARAMETERS::SHUFFLE_EPOCH, "Stateful actor and critic implies epoch shuffling");
+            if constexpr(PPO_SPEC::PARAMETERS::SHUFFLE_EPOCH){
+                auto observations_matrix = matrix_view(device, dataset.all_observations);
+                auto observations_privileged_matrix = matrix_view(device, dataset.all_observations_privileged);
+                for(TI dataset_i = 0; dataset_i < DATASET::STEPS_TOTAL; dataset_i++){
+                    TI sample_index = random::uniform_int_distribution(device.random, dataset_i, DATASET::STEPS_TOTAL-1, rng);
+                    {
+                        auto target_row = row(device, observations_matrix, dataset_i);
+                        auto source_row = row(device, observations_matrix, sample_index);
+                        swap(device, target_row, source_row);
+                    }
+                    {
+                        auto target_row = row(device, observations_privileged_matrix, dataset_i);
+                        auto source_row = row(device, observations_privileged_matrix, sample_index);
+                        swap(device, target_row, source_row);
+                    }
+                    if(PPO_SPEC::PARAMETERS::ADAPTIVE_LEARNING_RATE){
+                        auto target_row = row(device, dataset.actions_mean, dataset_i);
+                        auto source_row = row(device, dataset.actions_mean, sample_index);
+                        swap(device, target_row, source_row);
+                    }
+                    {
+                        auto target_row = row(device, dataset.actions, dataset_i);
+                        auto source_row = row(device, dataset.actions, sample_index);
+                        swap(device, target_row, source_row);
+                    }
+                    swap(device, dataset.advantages      , dataset.advantages      , dataset_i, 0, sample_index, 0);
+                    swap(device, dataset.action_log_probs, dataset.action_log_probs, dataset_i, 0, sample_index, 0);
+                    swap(device, dataset.target_values   , dataset.target_values   , dataset_i, 0, sample_index, 0);
+                }
+            }
+            static_assert(N_BATCHES > 0);
+            for(TI batch_i = 0; batch_i < N_BATCHES; batch_i++){
+                T batch_policy_kl_divergence = 0;
+                zero_gradient(device, ppo.critic);
+                zero_gradient(device, ppo.actor);
+
+                auto batch_offset = batch_i * BATCH_SIZE;
+                auto batch_observations            = view_range(device, dataset.all_observations              , batch_offset, tensor::ViewSpec<0, BATCH_SIZE>{});
+                auto batch_observations_privileged = view_range(device, dataset.all_observations_privileged   , batch_offset, tensor::ViewSpec<0, BATCH_SIZE>{});
+                auto batch_actions_mean            = view(device, dataset.actions_mean               , matrix::ViewSpec<BATCH_SIZE, ACTION_DIM                >(), batch_offset, 0);
+                auto batch_actions                 = view(device, dataset.actions                    , matrix::ViewSpec<BATCH_SIZE, ACTION_DIM                >(), batch_offset, 0);
+                auto batch_action_log_probs        = view(device, dataset.action_log_probs           , matrix::ViewSpec<BATCH_SIZE, 1                         >(), batch_offset, 0);
+                auto batch_advantages              = view(device, dataset.advantages                 , matrix::ViewSpec<BATCH_SIZE, 1                         >(), batch_offset, 0);
+                auto batch_target_values           = view(device, dataset.target_values              , matrix::ViewSpec<BATCH_SIZE, 1                         >(), batch_offset, 0);
+                auto batch_reset                   = view(device, dataset.reset                      , matrix::ViewSpec<BATCH_SIZE, 1                         >(), batch_offset, 0);
+
+                static constexpr TI STEPS = PPO_SPEC::PARAMETERS::STATEFUL_ACTOR_AND_CRITIC ? DATASET_SPEC::STEPS_PER_ENV : 1;
+                static constexpr TI FORWARD_BATCH_SIZE = PPO_SPEC::PARAMETERS::STATEFUL_ACTOR_AND_CRITIC ? DATASET_SPEC::SPEC::N_ENVIRONMENTS : BATCH_SIZE;
+                using ACTOR_INPUT_SHAPE = tensor::Prepend<tensor::Prepend<OBS_SHAPE, FORWARD_BATCH_SIZE>, STEPS>;
+                auto batch_observations_reshaped = reshape_row_major(device, batch_observations, ACTOR_INPUT_SHAPE{});
+                auto current_batch_actions_tensor = to_tensor(device, ppo_buffers.current_batch_actions);
+                auto current_batch_actions_tensor_reshaped = reshape_row_major(device, current_batch_actions_tensor, tensor::Shape<TI, STEPS, FORWARD_BATCH_SIZE, ACTION_DIM>{});
+                auto batch_reset_tensor_flat = to_tensor(device, batch_reset);
+                auto batch_reset_tensor = reshape_row_major(device, batch_reset_tensor_flat, tensor::Shape<TI, STEPS, FORWARD_BATCH_SIZE, 1>{});
+                Mode<mode::sequential::ResetMode<mode::Rollout<>, mode::sequential::ResetModeSpecification<TI, decltype(batch_reset_tensor)>>> mode;
+                mode.reset_container = batch_reset_tensor;
+                forward(device, ppo.actor, batch_observations_reshaped, current_batch_actions_tensor_reshaped, actor_buffers, rng, mode);
+
+                ppo_compute_actor_loss_gradient(device, ppo, ppo_buffers, batch_actions, batch_actions_mean, batch_action_log_probs, batch_advantages, policy_kl_divergence, batch_policy_kl_divergence, rng);
+                if(PPO_SPEC::PARAMETERS::ADAPTIVE_LEARNING_RATE){
+                    batch_policy_kl_divergence /= BATCH_SIZE;
+                    auto& actor_optimizer_parameters = get_ref(device, actor_optimizer.parameters, 0);
+                    if(batch_policy_kl_divergence > 2 * PPO_SPEC::PARAMETERS::ADAPTIVE_LEARNING_RATE_POLICY_KL_THRESHOLD){
+                        actor_optimizer_parameters.alpha = math::max(device.math, actor_optimizer_parameters.alpha * PPO_SPEC::PARAMETERS::ADAPTIVE_LEARNING_RATE_DECAY, PPO_SPEC::PARAMETERS::ADAPTIVE_LEARNING_RATE_MIN);
+                    }
+                    if(batch_policy_kl_divergence < 0.5 * PPO_SPEC::PARAMETERS::ADAPTIVE_LEARNING_RATE_POLICY_KL_THRESHOLD){
+                        actor_optimizer_parameters.alpha = math::min(device.math, actor_optimizer_parameters.alpha / PPO_SPEC::PARAMETERS::ADAPTIVE_LEARNING_RATE_DECAY, PPO_SPEC::PARAMETERS::ADAPTIVE_LEARNING_RATE_MAX);
+                    }
+                }
+                auto d_action_d_log_prob_action_tensor = to_tensor(device, ppo_buffers.d_action_log_prob_d_action);
+                auto d_action_d_log_prob_action_tensor_reshaped = reshape_row_major(device, d_action_d_log_prob_action_tensor, tensor::Shape<TI, STEPS, FORWARD_BATCH_SIZE, ACTION_DIM>{});
+                backward(device, ppo.actor, batch_observations_reshaped, d_action_d_log_prob_action_tensor_reshaped, actor_buffers, mode);
+
+                using CRITIC_INPUT_SHAPE = tensor::Prepend<tensor::Prepend<OBS_PRIV_SHAPE, FORWARD_BATCH_SIZE>, STEPS>;
+                auto batch_observations_privileged_reshaped = reshape_row_major(device, batch_observations_privileged, CRITIC_INPUT_SHAPE{});
+                {
+                    forward(device, ppo.critic, batch_observations_privileged_reshaped, critic_buffers, rng, mode);
+                    auto output_tensor = output(device, ppo.critic);
+                    static_assert(sizeof(output_tensor) <= sizeof(void*));
+                    auto output_matrix_view = matrix_view(device, output_tensor);
+                    nn::loss_functions::mse::gradient(device, output_matrix_view, batch_target_values, ppo_buffers.d_critic_output, 0.5);
+                    auto d_critic_output_tensor = to_tensor(device, ppo_buffers.d_critic_output);
+                    auto d_critic_output_tensor_reshaped = reshape_row_major(device, d_critic_output_tensor, tensor::Shape<TI, STEPS, FORWARD_BATCH_SIZE, 1>{});
+                    backward(device, ppo.critic, batch_observations_privileged_reshaped, d_critic_output_tensor_reshaped, critic_buffers, mode);
+                }
+                auto output_tensor = output(device, ppo.critic);
+                auto output_matrix_view = matrix_view(device, output_tensor);
+                T critic_loss = nn::loss_functions::mse::evaluate(device, output_matrix_view, batch_target_values);
+                add_scalar(device, device.logger, "ppo/critic_loss", critic_loss);
+                step(device, actor_optimizer, ppo.actor);
+                step(device, critic_optimizer, ppo.critic);
+            }
+        }
+        if(PPO_SPEC::PARAMETERS::ADAPTIVE_LEARNING_RATE) {
+            policy_kl_divergence /= N_EPOCHS * N_BATCHES * BATCH_SIZE;
+            add_scalar(device, device.logger, "ppo/policy_kl", policy_kl_divergence);
+        }
+    }
+
+    template <typename DEVICE_SOURCE, typename DEVICE_TARGET, typename PPO_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void copy(DEVICE_SOURCE& device_source, DEVICE_TARGET& device_target, const rl::algorithms::PPO<PPO_SPEC>& source, rl::algorithms::PPO<PPO_SPEC>& target){
+        copy(device_source, device_target, source.actor, target.actor);
+        copy(device_source, device_target, source.critic, target.critic);
+#ifdef RL_TOOLS_DEBUG_RL_ALGORITHMS_PPO_CHECK_INIT
+        target.initialized = source.initialized;
+#endif
+    }
+    template <typename DEVICE, typename SPEC_1, typename SPEC_2>
+    RL_TOOLS_FUNCTION_PLACEMENT typename SPEC_1::TYPE_POLICY::DEFAULT abs_diff(DEVICE& device, rl::algorithms::PPO<SPEC_1>& p1, rl::algorithms::PPO<SPEC_2>& p2){
+        using T = typename SPEC_1::TYPE_POLICY::DEFAULT;
+        T acc = 0;
+        acc += abs_diff(device, p1.actor, p2.actor);
+        acc += abs_diff(device, p1.critic, p2.critic);
+        return acc;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/algorithms/ppo/operations_generic_mixed_imitation.h

@@ -0,0 +1,81 @@
+#include "../../../version.h"
+#include "../../../rl_tools.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ALGORITHMS_PPO_OPERATIONS_GENERIC_MIXED_IMITATION_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ALGORITHMS_PPO_OPERATIONS_GENERIC_MIXED_IMITATION_H
+
+#include "mixed_imitation.h"
+#include "../../../containers/matrix/matrix.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT bool is_imitation_sample(DEVICE&, const SPEC&, typename SPEC::TI row){
+        return row % SPEC::GROUP_SIZE < SPEC::IMITATION_PER_GROUP;
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT bool use_teacher_action(DEVICE& device, const SPEC& spec, typename SPEC::TI row, typename SPEC::TI optimizer_updates){
+        return is_imitation_sample(device, spec, row) && optimizer_updates < SPEC::PARAMETERS::TEACHER_FORCING_UPDATES;
+    }
+
+    // The incoming derivative contains the RL subset mean, scaled for accumulation. IL rows
+    // replace it with half-MSE derivatives; the chosen weight is held constant in backward.
+    template <typename DEVICE, typename SPEC, typename PREDICTIONS_SPEC, typename TARGETS_SPEC, typename GRADIENT_SPEC, typename MODE>
+    RL_TOOLS_FUNCTION_PLACEMENT void mix_policy_output_gradients(DEVICE& device, const SPEC& spec, const Matrix<PREDICTIONS_SPEC>& predictions, const Matrix<TARGETS_SPEC>& targets, Matrix<GRADIENT_SPEC>& gradient, rl::algorithms::ppo::mixed_imitation::Metrics<SPEC>& metrics, typename SPEC::T accumulation_scale, Mode<MODE>){
+        using T = typename SPEC::T;
+        using TI = typename SPEC::TI;
+        using PARAMETERS = typename SPEC::PARAMETERS;
+        static_assert(SPEC::IMITATION_BATCH_SIZE > 0, "Mixed learning requires IL and RL in every minibatch");
+        static_assert(PREDICTIONS_SPEC::ROWS == SPEC::BATCH_SIZE && TARGETS_SPEC::ROWS == SPEC::BATCH_SIZE && GRADIENT_SPEC::ROWS == SPEC::BATCH_SIZE);
+        static_assert(PREDICTIONS_SPEC::COLS == TARGETS_SPEC::COLS && PREDICTIONS_SPEC::COLS == GRADIENT_SPEC::COLS);
+        static_assert(mode::is<MODE, rl::algorithms::ppo::mixed_imitation::FixedWeight> || mode::is<MODE, rl::algorithms::ppo::mixed_imitation::FixedNormRatio>);
+        constexpr TI ACTION_DIM = PREDICTIONS_SPEC::COLS;
+        constexpr T IL_ELEMENTS = static_cast<T>(SPEC::IMITATION_BATCH_SIZE * ACTION_DIM);
+        metrics = {};
+        T rl_squared_norm = 0;
+        T il_squared_norm = 0;
+        for(TI row = 0; row < SPEC::BATCH_SIZE; row++){
+            const bool imitation = is_imitation_sample(device, spec, row);
+            for(TI action = 0; action < ACTION_DIM; action++){
+                if(imitation){
+                    T error = static_cast<T>(get(predictions, row, action)) - static_cast<T>(get(targets, row, action));
+                    T derivative = error * accumulation_scale / IL_ELEMENTS;
+                    set(gradient, row, action, derivative);
+                    metrics.imitation_mse += error * error / IL_ELEMENTS;
+                    il_squared_norm += derivative * derivative;
+                }
+                else{
+                    T derivative = static_cast<T>(get(gradient, row, action));
+                    rl_squared_norm += derivative * derivative;
+                }
+            }
+        }
+        metrics.rl_output_gradient_norm = math::sqrt(device.math, rl_squared_norm);
+        metrics.imitation_output_gradient_norm = math::sqrt(device.math, il_squared_norm);
+        metrics.imitation_weight = PARAMETERS::IMITATION_WEIGHT;
+        if constexpr(mode::is<MODE, rl::algorithms::ppo::mixed_imitation::FixedNormRatio>){
+            if(metrics.rl_output_gradient_norm > PARAMETERS::NORM_EPSILON && metrics.imitation_output_gradient_norm > PARAMETERS::NORM_EPSILON){
+                T weight = PARAMETERS::OUTPUT_GRADIENT_NORM_RATIO * metrics.rl_output_gradient_norm / metrics.imitation_output_gradient_norm;
+                metrics.weight_clamped = weight > PARAMETERS::MAX_IMITATION_WEIGHT;
+                metrics.imitation_weight = math::min(device.math, weight, PARAMETERS::MAX_IMITATION_WEIGHT);
+            }
+            else{
+                metrics.balancing_fallback = true;
+            }
+        }
+        metrics.weighted_imitation_output_gradient_norm = metrics.imitation_weight * metrics.imitation_output_gradient_norm;
+        metrics.ratio_valid = metrics.rl_output_gradient_norm > PARAMETERS::NORM_EPSILON;
+        if(metrics.ratio_valid){
+            metrics.output_gradient_norm_ratio = metrics.weighted_imitation_output_gradient_norm / metrics.rl_output_gradient_norm;
+        }
+        for(TI row = 0; row < SPEC::BATCH_SIZE; row++){
+            if(is_imitation_sample(device, spec, row)){
+                for(TI action = 0; action < ACTION_DIM; action++){
+                    multiply(gradient, row, action, metrics.imitation_weight);
+                }
+            }
+        }
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/algorithms/ppo/persist.h

@@ -0,0 +1,28 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ALGORITHMS_PPO_PERSIST_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ALGORITHMS_PPO_PERSIST_H
+#include "ppo.h"
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template<typename DEVICE, typename SPEC, typename GROUP>
+    void save(DEVICE& device, rl::algorithms::PPO<SPEC>& ppo, GROUP& group) {
+        auto actor_group = create_group(device, group, "actor");
+        save(device, ppo.actor, actor_group);
+        auto critic_group = create_group(device, group, "critic");
+        save(device, ppo.critic, critic_group);
+    }
+    template<typename DEVICE, typename SPEC, typename GROUP>
+    bool load(DEVICE& device, rl::algorithms::PPO<SPEC>& ppo, GROUP& group) {
+        auto actor_group = get_group(device, group, "actor");
+        bool success = load(device, ppo.actor, actor_group);
+        auto critic_group = get_group(device, group, "critic");
+        success &= load(device, ppo.critic, critic_group);
+#ifdef RL_TOOLS_DEBUG_RL_ALGORITHMS_PPO_CHECK_INIT
+        ppo.initialized = true;
+#endif
+        return success;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/algorithms/ppo/ppo.h

@@ -0,0 +1,105 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ALGORITHMS_PPO_PPO_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ALGORITHMS_PPO_PPO_H
+
+#include "../../../rl/components/running_normalizer/running_normalizer.h"
+#include "../../../utils/generic/typing.h"
+#include "../../../numeric_types/categories.h"
+#include "../../components/on_policy_runner/collection.h"
+
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::rl::algorithms{
+    namespace ppo{
+        template<typename TYPE_POLICY, typename TI, TI T_BATCH_SIZE>
+        struct DefaultParameters {
+            using T = typename TYPE_POLICY::DEFAULT;
+            static constexpr T GAMMA = 0.99;
+            static constexpr T LAMBDA = 0.95;
+            static constexpr T EPSILON_CLIP = 0.2;
+            static constexpr T INITIAL_ACTION_STD = 0.5; // note this is NOT log(std) but actual std (log is applied at init)
+            static constexpr bool LEARN_ACTION_STD = true;
+            static constexpr T ACTION_ENTROPY_COEFFICIENT = 0.01;
+            static constexpr T ADVANTAGE_EPSILON = 1e-8;
+            static constexpr bool NORMALIZE_ADVANTAGE = true;
+            static constexpr bool ADAPTIVE_LEARNING_RATE = false;
+            static constexpr T ADAPTIVE_LEARNING_RATE_POLICY_KL_THRESHOLD = 0.008;
+            static constexpr T POLICY_KL_EPSILON = 1e-5;
+            static constexpr T ADAPTIVE_LEARNING_RATE_DECAY = (T)1/(T)1.5;
+            static constexpr T ADAPTIVE_LEARNING_RATE_MIN = 1e-6;
+            static constexpr T ADAPTIVE_LEARNING_RATE_MAX = 1e-2;
+            static constexpr TI N_WARMUP_STEPS_CRITIC = 0;
+            static constexpr TI N_WARMUP_STEPS_ACTOR = 0;
+            static constexpr TI N_EPOCHS = 10;
+            static constexpr TI BATCH_SIZE = T_BATCH_SIZE;
+            static constexpr bool IGNORE_TERMINATION = false; // ignoring the termination flag is useful for training on environments with negative rewards, where the agent would try to terminate the episode as soon as possible otherwise
+            static constexpr bool BOOTSTRAP_TRUNCATIONS = true;
+            static constexpr bool SHUFFLE_EPOCH = true;
+            static constexpr bool STATEFUL_ACTOR_AND_CRITIC = false;
+            static constexpr bool TRUNCATE_ON_EACH_ITERATION = false;
+        };
+
+        template<
+                typename T_TYPE_POLICY,
+                typename T_TI,
+                typename T_ENVIRONMENT,
+                typename T_ACTOR_TYPE,
+                typename T_CRITIC_TYPE,
+                typename T_PARAMETERS
+        >
+        struct Specification {
+            using TYPE_POLICY = T_TYPE_POLICY;
+            using TI = T_TI;
+            using ENVIRONMENT = T_ENVIRONMENT;
+            using ACTOR_TYPE = T_ACTOR_TYPE;
+            using CRITIC_TYPE = T_CRITIC_TYPE;
+            using PARAMETERS = T_PARAMETERS;
+            using COLLECTION_MODE = Mode<utils::typing::conditional_t<PARAMETERS::STATEFUL_ACTOR_AND_CRITIC,
+                mode::on_policy_runner::Sequential<mode::on_policy_runner::ActorCritic<>>,
+                mode::on_policy_runner::ActorCritic<>>>;
+            static constexpr bool ASYMMETRIC_OBSERVATIONS = !rl_tools::utils::typing::is_same_v<typename ENVIRONMENT::Observation, typename ENVIRONMENT::ObservationPrivileged>;
+
+            static_assert(get_last(typename ACTOR_TYPE::OUTPUT_SHAPE{}) == ENVIRONMENT::ACTION_DIM);
+            static_assert(get_last(typename CRITIC_TYPE::OUTPUT_SHAPE{}) == 1);
+        };
+
+        template <typename T_SPEC, bool T_DYNAMIC_ALLOCATION=true>
+        struct BufferSpecification{
+            using SPEC = T_SPEC;
+            static constexpr bool DYNAMIC_ALLOCATION = T_DYNAMIC_ALLOCATION;
+        };
+        template <typename T_BUFFER_SPEC>
+        struct Buffers{
+            using BUFFER_SPEC = T_BUFFER_SPEC;
+            using SPEC = typename BUFFER_SPEC::SPEC;
+            using T = typename SPEC::TYPE_POLICY::template GET<numeric_types::categories::Accumulator>;
+            using TI = typename SPEC::TI;
+            static constexpr TI BATCH_SIZE = SPEC::PARAMETERS::BATCH_SIZE;
+            static constexpr TI ACTION_DIM = SPEC::ENVIRONMENT::ACTION_DIM;
+            static constexpr TI OBSERVATION_DIM = SPEC::ENVIRONMENT::Observation::DIM;
+            Matrix<matrix::Specification<T, TI, BATCH_SIZE, ACTION_DIM, BUFFER_SPEC::DYNAMIC_ALLOCATION>> current_batch_actions;
+            Matrix<matrix::Specification<T, TI, BATCH_SIZE, 1, BUFFER_SPEC::DYNAMIC_ALLOCATION>> d_critic_output;
+            Matrix<matrix::Specification<T, TI, BATCH_SIZE, ACTION_DIM, BUFFER_SPEC::DYNAMIC_ALLOCATION>> d_action_log_prob_d_action;
+            Matrix<matrix::Specification<T, TI, BATCH_SIZE, ACTION_DIM, BUFFER_SPEC::DYNAMIC_ALLOCATION>> d_action_log_prob_d_action_log_std;
+            Matrix<matrix::Specification<T, TI, 1, ACTION_DIM/SPEC::ENVIRONMENT::N_AGENTS, BUFFER_SPEC::DYNAMIC_ALLOCATION>> rollout_log_std;
+        };
+    }
+
+    template<typename T_SPEC>
+    struct PPO {
+        using SPEC = T_SPEC;
+        using TYPE_POLICY = typename SPEC::TYPE_POLICY;
+        using TI = typename SPEC::TI;
+
+        typename SPEC::ACTOR_TYPE actor;
+        typename SPEC::CRITIC_TYPE critic;
+#ifdef RL_TOOLS_DEBUG_RL_ALGORITHMS_PPO_CHECK_INIT
+        bool initialized = false;
+#endif
+
+    };
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/algorithms/sac/loop/core/approximators_gru.h

@@ -0,0 +1,84 @@
+#include "../../../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ALGORITHMS_SAC_LOOP_CORE_APPROXIMATORS_GRU_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ALGORITHMS_SAC_LOOP_CORE_APPROXIMATORS_GRU_H
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::rl::algorithms::sac::loop::core{
+    template<typename TYPE_POLICY, typename TI, typename ENVIRONMENT, typename PARAMETERS, bool DYNAMIC_ALLOCATION=true>
+    struct ConfigApproximatorsGRU{
+        static constexpr bool USE_GRU = true;
+        using SAC_PARAMETERS = typename PARAMETERS::SAC_PARAMETERS;
+        template <typename CAPABILITY>
+        struct Actor{
+            using INPUT_SHAPE = tensor::Shape<TI, SAC_PARAMETERS::SEQUENCE_LENGTH, PARAMETERS::SAC_PARAMETERS::ACTOR_BATCH_SIZE, ENVIRONMENT::Observation::DIM>;
+            using INPUT_LAYER_CONFIG = nn::layers::dense::Configuration<TYPE_POLICY, TI, PARAMETERS::ACTOR_HIDDEN_DIM, PARAMETERS::CRITIC_ACTIVATION_FUNCTION, nn::layers::dense::DefaultInitializer<TYPE_POLICY, TI>, nn::parameters::groups::Input>;
+            using INPUT_LAYER = nn::layers::dense::BindConfiguration<INPUT_LAYER_CONFIG>;
+            using GRU_SPEC = nn::layers::gru::Configuration<TYPE_POLICY, TI, PARAMETERS::ACTOR_HIDDEN_DIM, nn::parameters::groups::Normal, true>;
+            using GRU = nn::layers::gru::BindConfiguration<GRU_SPEC>;
+            using GRU2_SPEC = nn::layers::gru::Configuration<TYPE_POLICY, TI, PARAMETERS::ACTOR_HIDDEN_DIM, nn::parameters::groups::Normal, true>;
+            using GRU2 = nn::layers::gru::BindConfiguration<GRU2_SPEC>;
+            using DENSE_LAYER_CONFIG = nn::layers::dense::Configuration<TYPE_POLICY, TI, PARAMETERS::ACTOR_HIDDEN_DIM, PARAMETERS::ACTOR_ACTIVATION_FUNCTION, nn::layers::dense::DefaultInitializer<TYPE_POLICY, TI>, nn::parameters::groups::Normal>;
+            using DENSE_LAYER = nn::layers::dense::BindConfiguration<DENSE_LAYER_CONFIG>;
+            using OUTPUT_LAYER_CONFIG = nn::layers::dense::Configuration<TYPE_POLICY, TI, 2*ENVIRONMENT::ACTION_DIM, nn::activation_functions::ActivationFunction::IDENTITY, nn::layers::dense::DefaultInitializer<TYPE_POLICY, TI>, nn::parameters::groups::Output>;
+            using OUTPUT_LAYER = nn::layers::dense::BindConfiguration<OUTPUT_LAYER_CONFIG>;
+            struct SAMPLE_AND_SQUASH_LAYER_PARAMETERS{
+                using T = typename TYPE_POLICY::DEFAULT;
+                static constexpr T LOG_STD_LOWER_BOUND = SAC_PARAMETERS::LOG_STD_LOWER_BOUND;
+                static constexpr T LOG_STD_UPPER_BOUND = SAC_PARAMETERS::LOG_STD_UPPER_BOUND;
+                static constexpr T LOG_PROBABILITY_EPSILON = SAC_PARAMETERS::LOG_PROBABILITY_EPSILON;
+                static constexpr bool ADAPTIVE_ALPHA = SAC_PARAMETERS::ADAPTIVE_ALPHA;
+                static constexpr bool UPDATE_ALPHA_WITH_ACTOR = false;
+                static constexpr T ALPHA = SAC_PARAMETERS::ALPHA;
+                static constexpr T TARGET_ENTROPY = SAC_PARAMETERS::TARGET_ENTROPY;
+            };
+            using SAMPLE_AND_SQUASH_LAYER_SPEC = nn::layers::sample_and_squash::Configuration<TYPE_POLICY, TI, SAMPLE_AND_SQUASH_LAYER_PARAMETERS>;
+            using SAMPLE_AND_SQUASH_LAYER = nn::layers::sample_and_squash::BindConfiguration<SAMPLE_AND_SQUASH_LAYER_SPEC>;
+            using SAMPLE_AND_SQUASH_MODULE = nn_models::sequential::Module<SAMPLE_AND_SQUASH_LAYER>;
+            using MODULE_GRU = nn_models::sequential::Module<INPUT_LAYER, nn_models::sequential::Module<GRU, nn_models::sequential::Module<OUTPUT_LAYER, SAMPLE_AND_SQUASH_MODULE>>>;
+            using MODULE_GRU_TWO_LAYER = nn_models::sequential::Module<INPUT_LAYER, nn_models::sequential::Module<GRU, nn_models::sequential::Module<GRU2, nn_models::sequential::Module<OUTPUT_LAYER, SAMPLE_AND_SQUASH_MODULE>>>>;
+            using MODULE_GRU_THREE_LAYER = nn_models::sequential::Module<INPUT_LAYER, nn_models::sequential::Module<GRU, nn_models::sequential::Module<GRU2, nn_models::sequential::Module<DENSE_LAYER, nn_models::sequential::Module<OUTPUT_LAYER, SAMPLE_AND_SQUASH_MODULE>>>>>;
+            using SELECTED_MODULE = rl_tools::utils::typing::conditional_t<PARAMETERS::CRITIC_NUM_LAYERS == 3, MODULE_GRU, rl_tools::utils::typing::conditional_t<PARAMETERS::CRITIC_NUM_LAYERS == 4, MODULE_GRU_TWO_LAYER, MODULE_GRU_THREE_LAYER>>;
+            static_assert(PARAMETERS::CRITIC_NUM_LAYERS == 3 || PARAMETERS::CRITIC_NUM_LAYERS == 4 || PARAMETERS::CRITIC_NUM_LAYERS == 5, "Only 3/4/5 layers (1 input + 1/2 GRU + 1/2 Output) are supported right now");
+            using MODEL = nn_models::sequential::Build<CAPABILITY, SELECTED_MODULE, INPUT_SHAPE>;
+        };
+        template <typename CAPABILITY>
+        struct Critic{
+            static constexpr TI INPUT_DIM = ENVIRONMENT::ObservationPrivileged::DIM+ENVIRONMENT::ACTION_DIM;
+            using INPUT_SHAPE = tensor::Shape<TI, SAC_PARAMETERS::SEQUENCE_LENGTH, PARAMETERS::SAC_PARAMETERS::CRITIC_BATCH_SIZE, INPUT_DIM>;
+            using INPUT_LAYER_CONFIG = nn::layers::dense::Configuration<TYPE_POLICY, TI, PARAMETERS::CRITIC_HIDDEN_DIM, PARAMETERS::CRITIC_ACTIVATION_FUNCTION, nn::layers::dense::DefaultInitializer<TYPE_POLICY, TI>, nn::parameters::groups::Input>;
+            using INPUT_LAYER = nn::layers::dense::BindConfiguration<INPUT_LAYER_CONFIG>;
+            using GRU_SPEC = nn::layers::gru::Configuration<TYPE_POLICY, TI, PARAMETERS::CRITIC_HIDDEN_DIM, nn::parameters::groups::Normal, true>;
+            using GRU = nn::layers::gru::BindConfiguration<GRU_SPEC>;
+            using GRU2_SPEC = nn::layers::gru::Configuration<TYPE_POLICY, TI, PARAMETERS::CRITIC_HIDDEN_DIM, nn::parameters::groups::Normal, true>;
+            using GRU2 = nn::layers::gru::BindConfiguration<GRU2_SPEC>;
+            using DENSE_LAYER_CONFIG = nn::layers::dense::Configuration<TYPE_POLICY, TI, PARAMETERS::CRITIC_HIDDEN_DIM, PARAMETERS::CRITIC_ACTIVATION_FUNCTION, nn::layers::dense::DefaultInitializer<TYPE_POLICY, TI>, nn::parameters::groups::Normal>;
+            using DENSE_LAYER = nn::layers::dense::BindConfiguration<DENSE_LAYER_CONFIG>;
+            using OUTPUT_LAYER_CONFIG = nn::layers::dense::Configuration<TYPE_POLICY, TI, 1, nn::activation_functions::ActivationFunction::IDENTITY, nn::layers::dense::DefaultInitializer<TYPE_POLICY, TI>, nn::parameters::groups::Output>;
+            using OUTPUT_LAYER = nn::layers::dense::BindConfiguration<OUTPUT_LAYER_CONFIG>;
+            using MODULE_GRU = nn_models::sequential::Module<INPUT_LAYER, nn_models::sequential::Module<GRU, nn_models::sequential::Module<OUTPUT_LAYER>>>;
+            using MODULE_GRU_TWO_LAYER = nn_models::sequential::Module<INPUT_LAYER, nn_models::sequential::Module<GRU, nn_models::sequential::Module<GRU2, nn_models::sequential::Module<OUTPUT_LAYER>>>>;
+            using MODULE_GRU_THREE_LAYER = nn_models::sequential::Module<INPUT_LAYER, nn_models::sequential::Module<GRU, nn_models::sequential::Module<GRU2, nn_models::sequential::Module<DENSE_LAYER, nn_models::sequential::Module<OUTPUT_LAYER>>>>>;
+            using SELECTED_MODULE = rl_tools::utils::typing::conditional_t<PARAMETERS::CRITIC_NUM_LAYERS == 3, MODULE_GRU, rl_tools::utils::typing::conditional_t<PARAMETERS::CRITIC_NUM_LAYERS == 4, MODULE_GRU_TWO_LAYER, MODULE_GRU_THREE_LAYER>>;
+            static_assert(PARAMETERS::CRITIC_NUM_LAYERS == 3 || PARAMETERS::CRITIC_NUM_LAYERS == 4 || PARAMETERS::CRITIC_NUM_LAYERS == 5, "Only 3/4/5 layers (1 input + 1/2 GRU + 1/2 Output) are supported right now");
+            using MODEL = nn_models::sequential::Build<CAPABILITY, SELECTED_MODULE, INPUT_SHAPE>;
+        };
+
+        using CAPABILITY_ACTOR = nn::capability::Gradient<nn::parameters::Adam, DYNAMIC_ALLOCATION>;
+        using CAPABILITY_CRITIC = nn::capability::Gradient<nn::parameters::Adam, DYNAMIC_ALLOCATION>;
+        using ACTOR_TYPE = typename Actor<CAPABILITY_ACTOR>::MODEL;
+        using CRITIC_TYPE = typename Critic<CAPABILITY_CRITIC>::MODEL;
+        using CRITIC_TARGET_TYPE = typename Critic<nn::capability::Forward<DYNAMIC_ALLOCATION>>::MODEL;
+        using ACTOR_OPTIMIZER_SPEC = nn::optimizers::adam::Specification<TYPE_POLICY, TI, typename PARAMETERS::ACTOR_OPTIMIZER_PARAMETERS, DYNAMIC_ALLOCATION>;
+        using CRITIC_OPTIMIZER_SPEC = nn::optimizers::adam::Specification<TYPE_POLICY, TI, typename PARAMETERS::CRITIC_OPTIMIZER_PARAMETERS, DYNAMIC_ALLOCATION>;
+        using ALPHA_OPTIMIZER_SPEC = nn::optimizers::adam::Specification<TYPE_POLICY, TI, typename PARAMETERS::ALPHA_OPTIMIZER_PARAMETERS, DYNAMIC_ALLOCATION>;
+        using ACTOR_OPTIMIZER = nn::optimizers::Adam<ACTOR_OPTIMIZER_SPEC>;
+        using CRITIC_OPTIMIZER = nn::optimizers::Adam<CRITIC_OPTIMIZER_SPEC>;
+        using ALPHA_OPTIMIZER = nn::optimizers::Adam<ALPHA_OPTIMIZER_SPEC>;
+
+    };
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif
+