hyperdrone 0.1.0__py3-none-any.whl

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

@@ -0,0 +1,62 @@
+#include "../../../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ALGORITHMS_PPO_LOOP_CORE_OPERATIONS_CUDA_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ALGORITHMS_PPO_LOOP_CORE_OPERATIONS_CUDA_H
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    namespace rl::algorithms::ppo::loop::core::cuda{
+        template <typename DEVICE, typename ENV_SPEC>
+        __global__
+        void init_environments_kernel(DEVICE device, Tensor<ENV_SPEC> envs){
+            using TI = typename DEVICE::index_t;
+            constexpr TI N = ENV_SPEC::SHAPE::template GET<0>;
+            TI env_i = threadIdx.x + blockIdx.x * blockDim.x;
+            if(env_i < N){
+                auto& env = get_ref(device, envs, env_i);
+                init(device, env);
+            }
+        }
+        template <typename DEVICE, typename ENV_SPEC>
+        __global__
+        void malloc_environments_kernel(DEVICE device, Tensor<ENV_SPEC> envs){
+            using TI = typename DEVICE::index_t;
+            constexpr TI N = ENV_SPEC::SHAPE::template GET<0>;
+            TI env_i = threadIdx.x + blockIdx.x * blockDim.x;
+            if(env_i < N){
+                auto& env = get_ref(device, envs, env_i);
+                malloc(device, env);
+            }
+        }
+    }
+    // CUDA overrides for the device-dependent helpers
+    template <typename DEV_SPEC, typename ENV_SPEC>
+    void malloc_environments(devices::CUDA<DEV_SPEC>& device, Tensor<ENV_SPEC>& envs){
+        using DEVICE = devices::CUDA<DEV_SPEC>;
+        using TI = typename DEVICE::index_t;
+        constexpr TI N = ENV_SPEC::SHAPE::template GET<0>;
+        constexpr TI BLOCKSIZE = 32;
+        constexpr TI N_BLOCKS = RL_TOOLS_DEVICES_CUDA_CEIL(N, BLOCKSIZE);
+        devices::cuda::TAG<DEVICE, true> tag_device{};
+        rl::algorithms::ppo::loop::core::cuda::malloc_environments_kernel<<<N_BLOCKS, BLOCKSIZE, 0, device.stream>>>(tag_device, envs);
+        check_status(device);
+    }
+    template <typename DEV_SPEC, typename ENV_SPEC>
+    void free_environments(devices::CUDA<DEV_SPEC>& device, Tensor<ENV_SPEC>& envs){
+        // L2F env free is a no-op, no kernel needed
+    }
+    template <typename DEV_SPEC, typename ENV_SPEC>
+    void init_environments(devices::CUDA<DEV_SPEC>& device, Tensor<ENV_SPEC>& envs){
+        using DEVICE = devices::CUDA<DEV_SPEC>;
+        using TI = typename DEVICE::index_t;
+        constexpr TI N = ENV_SPEC::SHAPE::template GET<0>;
+        constexpr TI BLOCKSIZE = 32;
+        constexpr TI N_BLOCKS = RL_TOOLS_DEVICES_CUDA_CEIL(N, BLOCKSIZE);
+        devices::cuda::TAG<DEVICE, true> tag_device{};
+        rl::algorithms::ppo::loop::core::cuda::init_environments_kernel<<<N_BLOCKS, BLOCKSIZE, 0, device.stream>>>(tag_device, envs);
+        check_status(device);
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#include "operations_generic.h"
+#endif

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

@@ -0,0 +1,245 @@
+#include "../../../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ALGORITHMS_PPO_LOOP_CORE_OPERATIONS_GENERIC_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ALGORITHMS_PPO_LOOP_CORE_OPERATIONS_GENERIC_H
+
+#include "../../../../../nn/optimizers/adam/instance/operations_generic.h"
+#include "../../../../../nn/layers/standardize/operations_generic.h"
+#include "../../../../../nn_models/mlp_unconditional_stddev/operations_generic.h"
+#include "../../../../../nn_models/sequential/operations_generic.h"
+#include "../../../../../nn/optimizers/adam/operations_generic.h"
+#include "../../../../../rl/algorithms/ppo/operations_generic.h"
+#include "../../../../../rl/components/on_policy_runner/operations_generic.h"
+#include "../../../../../random/operations_generic_array.h"
+
+#include "config.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    // Device-dependent helpers (overridden by operations_cuda.h for CUDA)
+    template <typename DEVICE, typename ENV_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc_environments(DEVICE& device, Tensor<ENV_SPEC>& envs){
+        using TI = typename DEVICE::index_t;
+        constexpr TI N = ENV_SPEC::SHAPE::template GET<0>;
+        for(TI env_i = 0; env_i < N; env_i++){
+            auto& env = get_ref(device, envs, env_i);
+            malloc(device, env);
+        }
+    }
+    template <typename DEVICE, typename ENV_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free_environments(DEVICE& device, Tensor<ENV_SPEC>& envs){
+        using TI = typename DEVICE::index_t;
+        constexpr TI N = ENV_SPEC::SHAPE::template GET<0>;
+        for(TI env_i = 0; env_i < N; env_i++){
+            auto& env = get_ref(device, envs, env_i);
+            free(device, env);
+        }
+    }
+    template <typename DEVICE, typename ENV_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void init_environments(DEVICE& device, Tensor<ENV_SPEC>& envs){
+        using TI = typename DEVICE::index_t;
+        constexpr TI N = ENV_SPEC::SHAPE::template GET<0>;
+        for(TI env_i = 0; env_i < N; env_i++){
+            auto& env = get_ref(device, envs, env_i);
+            init(device, env);
+        }
+    }
+    // Shared malloc/free/init
+    template <typename DEVICE, typename T_CONFIG>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, rl::algorithms::ppo::loop::core::State<T_CONFIG>& ts){
+        malloc(device, ts.rng);
+        malloc(device, ts.ppo);
+        malloc(device, ts.ppo_buffers);
+        malloc(device, ts.on_policy_runner_dataset);
+        malloc(device, ts.on_policy_runner);
+        malloc(device, ts.actor_eval_buffers);
+        malloc(device, ts.actor_buffers);
+        malloc(device, ts.critic_buffers);
+        malloc(device, ts.critic_buffers_gae);
+        malloc(device, ts.actor_optimizer);
+        malloc(device, ts.critic_optimizer);
+        malloc(device, ts.envs);
+        malloc(device, ts.env_parameters);
+        malloc_environments(device, ts.envs);
+    }
+    template <typename DEVICE, typename T_CONFIG>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, rl::algorithms::ppo::loop::core::State<T_CONFIG>& ts){
+        free(device, ts.rng);
+        free(device, ts.ppo);
+        free(device, ts.ppo_buffers);
+        free(device, ts.on_policy_runner_dataset);
+        free(device, ts.on_policy_runner);
+        free(device, ts.actor_eval_buffers);
+        free(device, ts.actor_buffers);
+        free(device, ts.critic_buffers);
+        free(device, ts.critic_buffers_gae);
+        free(device, ts.actor_optimizer);
+        free(device, ts.critic_optimizer);
+        free(device, ts.envs);
+        free(device, ts.env_parameters);
+        free_environments(device, ts.envs);
+    }
+    template <typename DEVICE, typename T_CONFIG>
+    RL_TOOLS_FUNCTION_PLACEMENT void init(DEVICE& device, rl::algorithms::ppo::loop::core::State<T_CONFIG>& ts, typename T_CONFIG::TI seed = 0){
+        using CONFIG = T_CONFIG;
+        using TI = typename DEVICE::index_t;
+
+        init(device, ts.rng, seed);
+        init_environments(device, ts.envs);
+        init(device, ts.ppo, ts.actor_optimizer, ts.critic_optimizer, ts.rng);
+        init(device, ts.on_policy_runner, ts.envs, ts.env_parameters, ts.ppo.actor, ts.rng);
+
+        ts.step = 0;
+    }
+    template <typename DEVICE, typename T_CONFIG>
+    RL_TOOLS_FUNCTION_PLACEMENT bool step(DEVICE& device, rl::algorithms::ppo::loop::core::State<T_CONFIG>& ts){
+        using CONFIG = T_CONFIG;
+        using TI = typename DEVICE::index_t;
+        constexpr TI CADENCE_PRE = CONFIG::CORE_PARAMETERS::STEP_LIMIT / 1000;
+        constexpr TI CADENCE = CADENCE_PRE > 0 ? CADENCE_PRE : 1;
+        constexpr TI N_AGENTS = T_CONFIG::ENVIRONMENT::N_AGENTS;
+        constexpr TI STEPS_TOTAL = CONFIG::ON_POLICY_RUNNER_DATASET_SPEC::STEPS_TOTAL;
+        set_step(device, device.logger, ts.step * CONFIG::CORE_PARAMETERS::N_ENVIRONMENTS * CONFIG::CORE_PARAMETERS::ON_POLICY_RUNNER_STEPS_PER_ENV);
+        bool finished = false;
+
+        if constexpr(T_CONFIG::CORE_PARAMETERS::NORMALIZE_OBSERVATIONS){
+            constexpr TI BATCH_SIZE = CONFIG::CORE_PARAMETERS::BATCH_SIZE;
+            constexpr TI N_BATCHES = STEPS_TOTAL / BATCH_SIZE;
+            static_assert(N_BATCHES > 0);
+            static constexpr TI ACCUMULATE_STEPS = CONFIG::PPO_SPEC::PARAMETERS::STATEFUL_ACTOR_AND_CRITIC ? CONFIG::ON_POLICY_RUNNER_DATASET_SPEC::STEPS_PER_ENV : 1;
+            static constexpr TI ACCUMULATE_FORWARD_BATCH_SIZE = CONFIG::PPO_SPEC::PARAMETERS::STATEFUL_ACTOR_AND_CRITIC ? CONFIG::ON_POLICY_RUNNER_DATASET_SPEC::SPEC::N_ENVIRONMENTS : BATCH_SIZE;
+            using OBS_SHAPE = typename T_CONFIG::ENVIRONMENT::Observation::SHAPE;
+            using OBS_PRIV_SHAPE = typename T_CONFIG::ENVIRONMENT::ObservationPrivileged::SHAPE;
+            using ACTOR_INPUT_SHAPE = tensor::Prepend<tensor::Prepend<OBS_SHAPE, ACCUMULATE_FORWARD_BATCH_SIZE>, ACCUMULATE_STEPS>;
+            using CRITIC_INPUT_SHAPE = tensor::Prepend<tensor::Prepend<OBS_PRIV_SHAPE, ACCUMULATE_FORWARD_BATCH_SIZE>, ACCUMULATE_STEPS>;
+            Mode<nn::layers::standardize::AccumulateMode<>> accumulate_mode;
+            if(ts.step == 0){
+                for(TI observation_normalization_warmup_step_i = 0; observation_normalization_warmup_step_i < T_CONFIG::OBSERVATION_NORMALIZATION_WARMUP_STEPS; observation_normalization_warmup_step_i++) {
+                    collect(device, ts.on_policy_runner_dataset, ts.on_policy_runner, ts.ppo.actor, ts.actor_eval_buffers, ts.rng);
+                    for(TI batch_i = 0; batch_i < N_BATCHES; batch_i++){
+                        auto batch_observations = view_range(device, ts.on_policy_runner_dataset.all_observations, batch_i * BATCH_SIZE, tensor::ViewSpec<0, BATCH_SIZE>{});
+                        auto batch_observations_reshaped = reshape_row_major(device, batch_observations, ACTOR_INPUT_SHAPE{});
+                        forward(device, ts.ppo.actor, batch_observations_reshaped, ts.actor_buffers, ts.rng, accumulate_mode);
+                        auto batch_observations_privileged = view_range(device, ts.on_policy_runner_dataset.all_observations_privileged, batch_i * BATCH_SIZE, tensor::ViewSpec<0, BATCH_SIZE>{});
+                        auto batch_observations_privileged_reshaped = reshape_row_major(device, batch_observations_privileged, CRITIC_INPUT_SHAPE{});
+                        forward(device, ts.ppo.critic, batch_observations_privileged_reshaped, ts.critic_buffers, ts.rng, accumulate_mode);
+                    }
+                }
+                init(device, ts.on_policy_runner, ts.envs, ts.env_parameters, ts.ppo.actor, ts.rng);
+            }
+        }
+        collect(device, ts.on_policy_runner_dataset, ts.on_policy_runner, ts.ppo.actor, ts.actor_eval_buffers, ts.rng);
+        if constexpr(T_CONFIG::CORE_PARAMETERS::NORMALIZE_OBSERVATIONS && T_CONFIG::CORE_PARAMETERS::NORMALIZE_OBSERVATIONS_CONTINUOUSLY){
+            constexpr TI BATCH_SIZE = CONFIG::CORE_PARAMETERS::BATCH_SIZE;
+            constexpr TI N_BATCHES = STEPS_TOTAL / BATCH_SIZE;
+            static constexpr TI ACCUMULATE_STEPS = CONFIG::PPO_SPEC::PARAMETERS::STATEFUL_ACTOR_AND_CRITIC ? CONFIG::ON_POLICY_RUNNER_DATASET_SPEC::STEPS_PER_ENV : 1;
+            static constexpr TI ACCUMULATE_FORWARD_BATCH_SIZE = CONFIG::PPO_SPEC::PARAMETERS::STATEFUL_ACTOR_AND_CRITIC ? CONFIG::ON_POLICY_RUNNER_DATASET_SPEC::SPEC::N_ENVIRONMENTS : BATCH_SIZE;
+            using OBS_SHAPE = typename T_CONFIG::ENVIRONMENT::Observation::SHAPE;
+            using OBS_PRIV_SHAPE = typename T_CONFIG::ENVIRONMENT::ObservationPrivileged::SHAPE;
+            using ACTOR_INPUT_SHAPE = tensor::Prepend<tensor::Prepend<OBS_SHAPE, ACCUMULATE_FORWARD_BATCH_SIZE>, ACCUMULATE_STEPS>;
+            using CRITIC_INPUT_SHAPE = tensor::Prepend<tensor::Prepend<OBS_PRIV_SHAPE, ACCUMULATE_FORWARD_BATCH_SIZE>, ACCUMULATE_STEPS>;
+            Mode<nn::layers::standardize::AccumulateMode<>> accumulate_mode;
+            for(TI batch_i = 0; batch_i < N_BATCHES; batch_i++){
+                auto batch_observations = view_range(device, ts.on_policy_runner_dataset.all_observations, batch_i * BATCH_SIZE, tensor::ViewSpec<0, BATCH_SIZE>{});
+                auto batch_observations_reshaped = reshape_row_major(device, batch_observations, ACTOR_INPUT_SHAPE{});
+                forward(device, ts.ppo.actor, batch_observations_reshaped, ts.actor_buffers, ts.rng, accumulate_mode);
+                auto batch_observations_privileged = view_range(device, ts.on_policy_runner_dataset.all_observations_privileged, batch_i * BATCH_SIZE, tensor::ViewSpec<0, BATCH_SIZE>{});
+                auto batch_observations_privileged_reshaped = reshape_row_major(device, batch_observations_privileged, CRITIC_INPUT_SHAPE{});
+                forward(device, ts.ppo.critic, batch_observations_privileged_reshaped, ts.critic_buffers, ts.rng, accumulate_mode);
+            }
+        }
+        static constexpr TI STEPS = CONFIG::PPO_SPEC::PARAMETERS::STATEFUL_ACTOR_AND_CRITIC ? CONFIG::ON_POLICY_RUNNER_DATASET_SPEC::STEPS_PER_ENV+1 : 1;
+        static constexpr TI FORWARD_BATCH_SIZE = CONFIG::PPO_SPEC::PARAMETERS::STATEFUL_ACTOR_AND_CRITIC ? CONFIG::ON_POLICY_RUNNER_DATASET_SPEC::SPEC::N_ENVIRONMENTS : CONFIG::ON_POLICY_RUNNER_DATASET_SPEC::STEPS_TOTAL_ALL;
+        using OBS_PRIV_SHAPE = typename CONFIG::ON_POLICY_RUNNER_DATASET_TYPE::OBS_PRIV_SHAPE;
+        using CRITIC_GAE_INPUT_SHAPE = tensor::Prepend<tensor::Prepend<OBS_PRIV_SHAPE, FORWARD_BATCH_SIZE>, STEPS>;
+        auto all_observations_privileged_reshaped = reshape_row_major(device, ts.on_policy_runner_dataset.all_observations_privileged, CRITIC_GAE_INPUT_SHAPE{});
+        auto all_values_tensor = to_tensor(device, ts.on_policy_runner_dataset.all_values);
+        auto all_values_tensor_reshaped = reshape_row_major(device, all_values_tensor, tensor::Shape<TI, STEPS, FORWARD_BATCH_SIZE, 1>{});
+        auto all_reset_tensor = to_tensor(device, ts.on_policy_runner_dataset.all_reset);
+        auto all_reset_tensor_reshaped = reshape_row_major(device, all_reset_tensor, tensor::Shape<TI, STEPS, FORWARD_BATCH_SIZE, 1>{});
+        Mode<nn::layers::gru::ResetMode<mode::Rollout<>, nn::layers::gru::ResetModeSpecification<TI, decltype(all_reset_tensor_reshaped)>>> critic_reset_mode;
+        critic_reset_mode.reset_container = all_reset_tensor_reshaped;
+        evaluate(device, ts.ppo.critic, all_observations_privileged_reshaped, all_values_tensor_reshaped, ts.critic_buffers_gae, ts.rng, critic_reset_mode);
+        estimate_generalized_advantages(device, ts.on_policy_runner_dataset, typename CONFIG::PPO_TYPE::SPEC::PARAMETERS{});
+        train(device, ts.ppo, ts.on_policy_runner_dataset, ts.actor_optimizer, ts.critic_optimizer, ts.ppo_buffers, ts.actor_buffers, ts.critic_buffers, ts.rng);
+
+#ifndef RL_TOOLS_DISABLE_TENSORBOARD
+        {
+            // logging actor std
+            auto& last_layer = get_last_layer(ts.ppo.actor);
+            using T = typename CONFIG::T;
+            constexpr TI PER_AGENT_ACTION_DIM = T_CONFIG::ENVIRONMENT::ACTION_DIM/N_AGENTS;
+            for(TI action_i = 0; action_i < PER_AGENT_ACTION_DIM; action_i++){
+                T current_action_log_std = get(device, last_layer.log_std.parameters, action_i % PER_AGENT_ACTION_DIM);
+                add_scalar(device, device.logger, "actor/log_std", current_action_log_std, CADENCE);
+            }
+        }
+#endif
+
+        add_scalar(device, device.logger, "ppo/step", ts.step, CADENCE);
+
+        ts.step++;
+        if(ts.step > CONFIG::CORE_PARAMETERS::STEP_LIMIT){
+            return true;
+        }
+        else{
+            return finished;
+        }
+    }
+    template <typename DEVICE, typename PARAMETERS, typename utils::typing::enable_if<utils::typing::is_same_v<typename PARAMETERS::TAG, rl::algorithms::ppo::loop::core::ParametersTag>>::type* = nullptr>
+    RL_TOOLS_FUNCTION_PLACEMENT void log(DEVICE& device, PARAMETERS){
+        log(device, device.logger, "STEP_LIMIT: ", PARAMETERS::STEP_LIMIT);
+        log(device, device.logger, "ACTOR_HIDDEN_DIM: ", PARAMETERS::ACTOR_HIDDEN_DIM);
+        log(device, device.logger, "ACTOR_NUM_LAYERS: ", PARAMETERS::ACTOR_NUM_LAYERS);
+        log(device, device.logger, "ACTOR_ACTIVATION_FUNCTION: ", PARAMETERS::ACTOR_ACTIVATION_FUNCTION);
+        log(device, device.logger, "CRITIC_HIDDEN_DIM: ", PARAMETERS::CRITIC_HIDDEN_DIM);
+        log(device, device.logger, "CRITIC_NUM_LAYERS: ", PARAMETERS::CRITIC_NUM_LAYERS);
+        log(device, device.logger, "CRITIC_ACTIVATION_FUNCTION: ", PARAMETERS::CRITIC_ACTIVATION_FUNCTION);
+        log(device, device.logger, "EPISODE_STEP_LIMIT: ", PARAMETERS::EPISODE_STEP_LIMIT);
+        log(device, device.logger, "N_ENVIRONMENTS: ", PARAMETERS::N_ENVIRONMENTS);
+        log(device, device.logger, "ON_POLICY_RUNNER_STEPS_PER_ENV: ", PARAMETERS::ON_POLICY_RUNNER_STEPS_PER_ENV);
+        log(device, device.logger, "BATCH_SIZE: ", PARAMETERS::BATCH_SIZE);
+    }
+    template <typename DEVICE, typename CONFIG, typename utils::typing::enable_if<utils::typing::is_same_v<typename CONFIG::TAG, rl::algorithms::ppo::loop::core::ConfigTag>>::type* = nullptr>
+    RL_TOOLS_FUNCTION_PLACEMENT void log(DEVICE& device, CONFIG){
+        log(device, typename CONFIG::CORE_PARAMETERS{});
+    }
+
+    template <auto INDEX, typename DEVICE, typename T_CONFIG>
+    constexpr auto& get_nn(DEVICE& device, rl::algorithms::ppo::loop::core::State<T_CONFIG>& ts){
+        static_assert(INDEX < T_CONFIG::NUM_NNS, "Index out of bounds, there are only 2 neural networks in the PPO");
+        if constexpr(INDEX == 0){
+            return ts.ppo.actor;
+        }
+        else{
+            return ts.ppo.critic;
+        }
+    }
+    template <auto INDEX, typename DEVICE, typename T_CONFIG>
+    constexpr auto& get_nn_name(DEVICE& device, rl::algorithms::ppo::loop::core::State<T_CONFIG>& ts){
+        static_assert(INDEX < T_CONFIG::NUM_NNS, "Index out of bounds, there are only 2 neural networks in the PPO");
+        if constexpr(INDEX == 0){
+            return "actor";
+        }
+        else{
+            return "critic";
+        }
+    }
+    template <typename DEVICE, typename CONFIG>
+    typename CONFIG::T abs_diff(DEVICE& device, rl::algorithms::ppo::loop::core::State<CONFIG>& s1, rl::algorithms::ppo::loop::core::State<CONFIG>& s2){
+        using T = typename CONFIG::T;
+        T acc = 0;
+        acc += abs_diff(device, s1.actor_optimizer, s2.actor_optimizer);
+        acc += abs_diff(device, s1.critic_optimizer, s2.critic_optimizer);
+        acc += abs_diff(device, s1.ppo, s2.ppo);
+        acc += abs_diff(device, s1.on_policy_runner, s2.on_policy_runner);
+        acc += abs_diff(device, s1.on_policy_runner_dataset, s2.on_policy_runner_dataset);
+        acc += math::abs(device.math, (T)s1.step - (T)s2.step);
+        acc += math::abs(device.math, (T)s1.next_checkpoint_id - (T)s2.next_checkpoint_id);
+        acc += math::abs(device.math, (T)s1.next_evaluation_id - (T)s2.next_evaluation_id);
+        acc += abs_diff(device, s1.rng, s2.rng);
+        return acc;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

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

@@ -0,0 +1,102 @@
+#include "../../../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ALGORITHMS_PPO_LOOP_CORE_PERSIST_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ALGORITHMS_PPO_LOOP_CORE_PERSIST_H
+#include "state.h"
+#include "../../../../../nn/optimizers/adam/persist.h"
+#include "../../../../../nn/optimizers/adam/instance/persist.h"
+#include "../../../../../nn_models/mlp_unconditional_stddev/persist.h"  // Must be before sequential for correct overload resolution
+#include "../../../../../nn_models/sequential/persist.h"
+#include "../../../../../rl/algorithms/ppo/persist.h"
+#include "../../../../../rl/components/on_policy_runner/persist.h"
+#include "../../../../../random/persist.h"
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template <typename DEVICE, typename T_CONFIG, typename GROUP>
+    void save(DEVICE& device, rl::algorithms::ppo::loop::core::State<T_CONFIG>& ts, GROUP& group){
+        using TI = typename DEVICE::index_t;
+        using T = typename T_CONFIG::T;
+        auto actor_optimizer_group = create_group(device, group, "actor_optimizer");
+        save(device, ts.actor_optimizer, actor_optimizer_group);
+        auto critic_optimizer_group = create_group(device, group, "critic_optimizer");
+        save(device, ts.critic_optimizer, critic_optimizer_group);
+        auto ppo_group = create_group(device, group, "ppo");
+        save(device, ts.ppo, ppo_group);
+        auto on_policy_runner_group = create_group(device, group, "on_policy_runner");
+        save(device, ts.on_policy_runner, on_policy_runner_group);
+        auto on_policy_runner_dataset_group = create_group(device, group, "on_policy_runner_dataset");
+        save(device, ts.on_policy_runner_dataset, on_policy_runner_dataset_group);
+        auto rng_group = create_group(device, group, "rng");
+        save(device, ts.rng, rng_group);
+        Tensor<tensor::Specification<TI, TI, tensor::Shape<TI, 1>>> step_tensor;
+        malloc(device, step_tensor);
+        set(device, step_tensor, ts.step, 0);
+        save(device, step_tensor, group, "step");
+        free(device, step_tensor);
+        Tensor<tensor::Specification<TI, TI, tensor::Shape<TI, 1>>> next_checkpoint_id_tensor;
+        malloc(device, next_checkpoint_id_tensor);
+        set(device, next_checkpoint_id_tensor, ts.next_checkpoint_id, 0);
+        save(device, next_checkpoint_id_tensor, group, "next_checkpoint_id");
+        free(device, next_checkpoint_id_tensor);
+        Tensor<tensor::Specification<TI, TI, tensor::Shape<TI, 1>>> next_evaluation_id_tensor;
+        malloc(device, next_evaluation_id_tensor);
+        set(device, next_evaluation_id_tensor, ts.next_evaluation_id, 0);
+        save(device, next_evaluation_id_tensor, group, "next_evaluation_id");
+        free(device, next_evaluation_id_tensor);
+    }
+    template <typename DEVICE, typename T_CONFIG, typename GROUP>
+    bool load(DEVICE& device, rl::algorithms::ppo::loop::core::State<T_CONFIG>& ts, GROUP& group){
+        using TI = typename DEVICE::index_t;
+        using T = typename T_CONFIG::T;
+        bool success = true;
+        bool step_result;
+        auto actor_optimizer_group = get_group(device, group, "actor_optimizer");
+        step_result = load(device, ts.actor_optimizer, actor_optimizer_group);
+        if(!step_result){ log(device, device.logger, "PPO loop load failed: actor_optimizer"); }
+        success &= step_result;
+        auto critic_optimizer_group = get_group(device, group, "critic_optimizer");
+        step_result = load(device, ts.critic_optimizer, critic_optimizer_group);
+        if(!step_result){ log(device, device.logger, "PPO loop load failed: critic_optimizer"); }
+        success &= step_result;
+        auto ppo_group = get_group(device, group, "ppo");
+        step_result = load(device, ts.ppo, ppo_group);
+        if(!step_result){ log(device, device.logger, "PPO loop load failed: ppo"); }
+        success &= step_result;
+        auto on_policy_runner_group = get_group(device, group, "on_policy_runner");
+        step_result = load(device, ts.on_policy_runner, on_policy_runner_group);
+        if(!step_result){ log(device, device.logger, "PPO loop load failed: on_policy_runner"); }
+        success &= step_result;
+        auto on_policy_runner_dataset_group = get_group(device, group, "on_policy_runner_dataset");
+        step_result = load(device, ts.on_policy_runner_dataset, on_policy_runner_dataset_group);
+        if(!step_result){ log(device, device.logger, "PPO loop load failed: on_policy_runner_dataset"); }
+        success &= step_result;
+        auto rng_group = get_group(device, group, "rng");
+        step_result = load(device, ts.rng, rng_group);
+        if(!step_result){ log(device, device.logger, "PPO loop load failed: rng"); }
+        success &= step_result;
+        Tensor<tensor::Specification<TI, TI, tensor::Shape<TI, 1>>> step_tensor;
+        malloc(device, step_tensor);
+        step_result = load(device, step_tensor, group, "step");
+        if(!step_result){ log(device, device.logger, "PPO loop load failed: step"); }
+        success &= step_result;
+        ts.step = get(device, step_tensor, 0);
+        free(device, step_tensor);
+        Tensor<tensor::Specification<TI, TI, tensor::Shape<TI, 1>>> next_checkpoint_id_tensor;
+        malloc(device, next_checkpoint_id_tensor);
+        step_result = load(device, next_checkpoint_id_tensor, group, "next_checkpoint_id");
+        if(!step_result){ log(device, device.logger, "PPO loop load failed: next_checkpoint_id"); }
+        success &= step_result;
+        ts.next_checkpoint_id = get(device, next_checkpoint_id_tensor, 0);
+        free(device, next_checkpoint_id_tensor);
+        Tensor<tensor::Specification<TI, TI, tensor::Shape<TI, 1>>> next_evaluation_id_tensor;
+        malloc(device, next_evaluation_id_tensor);
+        step_result = load(device, next_evaluation_id_tensor, group, "next_evaluation_id");
+        if(!step_result){ log(device, device.logger, "PPO loop load failed: next_evaluation_id"); }
+        success &= step_result;
+        ts.next_evaluation_id = get(device, next_evaluation_id_tensor, 0);
+        free(device, next_evaluation_id_tensor);
+        return success;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

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

@@ -0,0 +1,52 @@
+#include "../../../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ALGORITHMS_PPO_LOOP_CORE_STATE_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ALGORITHMS_PPO_LOOP_CORE_STATE_H
+
+#include "../../../../../rl/algorithms/ppo/ppo.h"
+#include "../../../../../rl/components/off_policy_runner/off_policy_runner.h"
+#include "../../../../../rl/environments/environments.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    namespace rl::algorithms::ppo::loop::core{
+        // Config State (Init/Step)
+        template<typename T_CONFIG>
+        struct State{
+            using CONFIG = T_CONFIG;
+            using TYPE_POLICY = typename CONFIG::TYPE_POLICY;
+            using TI = typename CONFIG::TI;
+            using T = typename TYPE_POLICY::DEFAULT;
+            static constexpr bool DYNAMIC_ALLOCATION = CONFIG::DYNAMIC_ALLOCATION;
+            typename CONFIG::NN::ACTOR_OPTIMIZER actor_optimizer;
+            typename CONFIG::NN::CRITIC_OPTIMIZER critic_optimizer;
+            typename CONFIG::RNG rng;
+            typename CONFIG::PPO_TYPE ppo;
+            typename CONFIG::PPO_BUFFERS_TYPE ppo_buffers;
+            typename CONFIG::ON_POLICY_RUNNER_TYPE on_policy_runner;
+            typename CONFIG::ON_POLICY_RUNNER_DATASET_TYPE on_policy_runner_dataset;
+            typename CONFIG::ACTOR_EVAL_BUFFERS actor_eval_buffers;
+            typename CONFIG::ACTOR_BUFFERS actor_buffers;
+            typename CONFIG::CRITIC_BUFFERS critic_buffers;
+            typename CONFIG::CRITIC_BUFFERS_GAE critic_buffers_gae;
+            Tensor<tensor::Specification<typename CONFIG::ENVIRONMENT, TI, tensor::Shape<TI, CONFIG::CORE_PARAMETERS::N_ENVIRONMENTS>, DYNAMIC_ALLOCATION>> envs;
+            Tensor<tensor::Specification<typename CONFIG::ENVIRONMENT::Parameters, TI, tensor::Shape<TI, CONFIG::CORE_PARAMETERS::N_ENVIRONMENTS>, DYNAMIC_ALLOCATION>> env_parameters;
+//            MatrixDynamic<matrix::Specification<typename CONFIG::T, TI, 1, CONFIG::ENVIRONMENT::Observation::DIM>> observations_mean, observations_std;
+            environments::DummyUI ui;
+            TI next_checkpoint_id = 0;
+            TI next_evaluation_id = 0;
+            TI step;
+        };
+    }
+    template <typename T_CONFIG>
+    constexpr auto& get_actor(rl::algorithms::ppo::loop::core::State<T_CONFIG>& ts){
+        return ts.ppo.actor;
+    }
+    template <typename T_CONFIG>
+    constexpr auto& get_actor(const rl::algorithms::ppo::loop::core::State<T_CONFIG>& ts){
+        return ts.ppo.actor;
+    }
+
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif