hyperdrone 0.1.0__py3-none-any.whl

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/components/on_policy_runner/operations_generic_extensions.h

@@ -0,0 +1,98 @@
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    namespace rl::components::on_policy_runner{
+        template <typename T_SPEC>
+        struct CollectionEvaluationBuffer{
+            using SPEC = T_SPEC;
+            using T = typename SPEC::T;
+            using TI = typename SPEC::TI;
+            typename SPEC::CONTAINER_TYPE_TAG::template type<matrix::Specification<T, TI, SPEC::N_ENVIRONMENTS, SPEC::ENVIRONMENT::Observation::DIM>> observations;
+            typename SPEC::CONTAINER_TYPE_TAG::template type<matrix::Specification<T, TI, SPEC::N_ENVIRONMENTS, SPEC::ENVIRONMENT::ACTION_DIM>> actions;
+        };
+    }
+    template <typename DEVICE, typename SPEC>
+    void malloc(DEVICE& device, rl::components::on_policy_runner::CollectionEvaluationBuffer<SPEC>& buffer){
+        malloc(device, buffer.observations);
+        malloc(device, buffer.actions);
+    }
+    template <typename DEVICE, typename SPEC>
+    void free(DEVICE& device, rl::components::on_policy_runner::CollectionEvaluationBuffer<SPEC>& buffer){
+        free(device, buffer.observations);
+        free(device, buffer.actions);
+    }
+    template <typename DEVICE, typename DEVICE_EVALUATION, typename DATASET_SPEC, typename ACTOR, typename ACTOR_EVALUATION, typename RNG, typename RNG_EVAL> // todo: make this not PPO but general policy with output distribution
+    void collect_hybrid(DEVICE& device, DEVICE_EVALUATION& device_evaluation, rl::components::on_policy_runner::Dataset<DATASET_SPEC>& dataset, rl::components::OnPolicyRunner<typename DATASET_SPEC::SPEC>& runner, ACTOR& actor, ACTOR_EVALUATION& actor_evaluation, typename ACTOR_EVALUATION::template Buffer<DATASET_SPEC::SPEC::N_ENVIRONMENTS>& policy_eval_buffers, rl::components::on_policy_runner::CollectionEvaluationBuffer<typename DATASET_SPEC::SPEC> evaluation_buffer, rl::components::on_policy_runner::CollectionEvaluationBuffer<typename DATASET_SPEC::SPEC>& evaluation_buffer_evaluation, RNG& rng, RNG_EVAL& rng_evaluation){
+#ifdef RL_TOOLS_DEBUG_RL_COMPONENTS_ON_POLICY_RUNNER_CHECK_INIT
+        utils::assert_exit(device, runner.initialized, "rl::components::on_policy_runner::collect: runner not initialized");
+#endif
+        using SPEC = typename DATASET_SPEC::SPEC;
+        using BUFFER = rl::components::on_policy_runner::Dataset<SPEC>;
+        using T = typename SPEC::T;
+        using TI = typename SPEC::TI;
+//        TI prologue_time = 0;
+//        TI copy_observations_time = 0;
+//        TI evaluate_time = 0;
+//        TI copy_back_time = 0;
+//        TI epilogue_time = 0;
+        for(TI step_i = 0; step_i < DATASET_SPEC::STEPS_PER_ENV; step_i++){
+            auto actions_mean            = view(device, dataset.actions_mean               , matrix::ViewSpec<SPEC::N_ENVIRONMENTS, SPEC::ENVIRONMENT::ACTION_DIM>()                , step_i*SPEC::N_ENVIRONMENTS, 0);
+            auto actions                 = view(device, dataset.actions                    , matrix::ViewSpec<SPEC::N_ENVIRONMENTS, SPEC::ENVIRONMENT::ACTION_DIM>()                , step_i*SPEC::N_ENVIRONMENTS, 0);
+            auto observations            = view(device, dataset.observations               , matrix::ViewSpec<SPEC::N_ENVIRONMENTS, SPEC::ENVIRONMENT::Observation::DIM>()          , step_i*SPEC::N_ENVIRONMENTS, 0);
+            auto observations_privileged = view(device, dataset.all_observations_privileged, matrix::ViewSpec<SPEC::N_ENVIRONMENTS, SPEC::ENVIRONMENT::ObservationPrivileged::DIM>(), step_i*SPEC::N_ENVIRONMENTS, 0);
+
+            {
+//                auto start = std::chrono::high_resolution_clock::now();
+                rl::components::on_policy_runner::prologue(device, observations_privileged, observations, runner, rng, step_i);
+//                auto end = std::chrono::high_resolution_clock::now();
+//                prologue_time += std::chrono::duration_cast<std::chrono::microseconds>(end - start).count();
+            }
+            {
+//                auto start = std::chrono::high_resolution_clock::now();
+                copy(device, device_evaluation, observations, evaluation_buffer_evaluation.observations);
+//                auto end = std::chrono::high_resolution_clock::now();
+//                copy_observations_time += std::chrono::duration_cast<std::chrono::microseconds>(end - start).count();
+            }
+            {
+//                auto start = std::chrono::high_resolution_clock::now();
+                evaluate(device_evaluation, actor_evaluation, evaluation_buffer_evaluation.observations, evaluation_buffer_evaluation.actions, policy_eval_buffers, rng, rng_evaluation);
+                cudaDeviceSynchronize();
+//                auto end = std::chrono::high_resolution_clock::now();
+//                evaluate_time += std::chrono::duration_cast<std::chrono::microseconds>(end - start).count();
+            }
+            {
+
+//                auto start = std::chrono::high_resolution_clock::now();
+                copy(device_evaluation, device, evaluation_buffer_evaluation.actions, evaluation_buffer.actions);
+                copy(device, device, evaluation_buffer.actions, actions_mean);
+//                auto end = std::chrono::high_resolution_clock::now();
+//                copy_back_time += std::chrono::duration_cast<std::chrono::microseconds>(end - start).count();
+            }
+            {
+//                auto start = std::chrono::high_resolution_clock::now();
+                auto& last_layer = get_last_layer(actor);
+                rl::components::on_policy_runner::epilogue(device, dataset, runner, actions_mean, actions, last_layer.log_std.parameters, rng, step_i);
+//                auto end = std::chrono::high_resolution_clock::now();
+//                epilogue_time += std::chrono::duration_cast<std::chrono::microseconds>(end - start).count();
+            }
+        }
+//        std::cout << "prologue_time: " << prologue_time << std::endl;
+//        std::cout << "copy_observations_time: " << copy_observations_time << std::endl;
+//        std::cout << "evaluate_time: " << evaluate_time << std::endl;
+//        std::cout << "copy_back_time: " << copy_back_time << std::endl;
+//        std::cout << "epilogue_time: " << epilogue_time << std::endl;
+
+        // final observation
+        for(TI env_i = 0; env_i < SPEC::N_ENVIRONMENTS; env_i++){
+            auto& env = get(runner.environments, 0, env_i);
+            auto& env_parameters = get(runner.env_parameters, 0, env_i);
+            auto& state = get(runner.states, 0, env_i);
+            TI row_i = DATASET_SPEC::STEPS_PER_ENV * SPEC::N_ENVIRONMENTS + env_i;
+            auto observation = row(device, dataset.all_observations_privileged, row_i);
+            observe(device, env, env_parameters, state, typename SPEC::ENVIRONMENT::ObservationPrivileged{}, observation, rng);
+//            auto observation = row(device, dataset.all_observations_normalized, row_i);
+//            normalize(device, observations_mean, observations_std, observation, observation_normalized);
+        }
+        runner.step += SPEC::N_ENVIRONMENTS * DATASET_SPEC::STEPS_PER_ENV;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/components/on_policy_runner/operations_generic_per_env.h

@@ -0,0 +1,91 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_COMPONENTS_ON_POLICY_RUNNER_OPERATIONS_GENERIC_PER_ENV_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_COMPONENTS_ON_POLICY_RUNNER_OPERATIONS_GENERIC_PER_ENV_H
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::rl::components::on_policy_runner::per_env{
+    template <typename DEVICE, typename OBS_PRIV_SPEC, typename OBS_SPEC, typename SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT void prologue(DEVICE& device, Tensor<OBS_PRIV_SPEC>& observations_privileged, Tensor<OBS_SPEC>& observations, rl::components::OnPolicyRunner<SPEC>& runner, RNG& rng, typename DEVICE::index_t env_i){
+        auto& env = get(runner.environments, 0, env_i);
+        auto& state = get(runner.states, 0, env_i);
+        auto& parameters = get(runner.env_parameters, 0, env_i);
+        if(get(runner.truncated, 0, env_i)){
+#ifndef __CUDA_ARCH__
+            if (env_i < SPEC::EPISODE_STATS_N_ENVIRONMENTS){
+                add_scalar(device, device.logger, "episode/length", get(runner.episode_step, 0, env_i), SPEC::EPISODE_STATS_CADENCE);
+                add_scalar(device, device.logger, "episode/return", get(runner.episode_return, 0, env_i), SPEC::EPISODE_STATS_CADENCE);
+            }
+#endif
+            set(runner.truncated, 0, env_i, false);
+            set(runner.episode_step, 0, env_i, 0);
+            set(runner.episode_return, 0, env_i, 0);
+            sample_initial_parameters(device, env, parameters, rng);
+            sample_initial_state(device, env, parameters, state, rng);
+        }
+        auto obs_slice = view(device, observations, env_i);
+        auto obs_flat = view_memory<tensor::Shape<typename DEVICE::index_t, SPEC::ENVIRONMENT::Observation::DIM>>(device, obs_slice);
+        auto obs_matrix = matrix_view(device, obs_flat);
+        observe(device, env, parameters, state, typename SPEC::ENVIRONMENT::Observation{}, obs_matrix, rng);
+
+        auto obs_priv_slice = view(device, observations_privileged, env_i);
+        if constexpr(!rl_tools::utils::typing::is_same_v<typename SPEC::ENVIRONMENT::Observation, typename SPEC::ENVIRONMENT::ObservationPrivileged>){
+            auto obs_priv_flat = view_memory<tensor::Shape<typename DEVICE::index_t, SPEC::ENVIRONMENT::ObservationPrivileged::DIM>>(device, obs_priv_slice);
+            auto obs_priv_matrix = matrix_view(device, obs_priv_flat);
+            observe(device, env, parameters, state, typename SPEC::ENVIRONMENT::ObservationPrivileged{}, obs_priv_matrix, rng);
+        }
+        else{
+            copy(device, device, obs_slice, obs_priv_slice);
+        }
+    }
+    template <typename DEVICE, typename DATASET_SPEC, typename ACTIONS_MEAN_SPEC, typename ACTIONS_SPEC, typename ACTION_LOG_STD_SPEC, typename RNG> // todo: make this not PPO but general policy with output distribution
+    RL_TOOLS_FUNCTION_PLACEMENT void epilogue(DEVICE& device, rl::components::on_policy_runner::Dataset<DATASET_SPEC>& dataset, rl::components::OnPolicyRunner<typename DATASET_SPEC::SPEC>& runner, Matrix<ACTIONS_MEAN_SPEC>& actions_mean, Matrix<ACTIONS_SPEC>& actions, Matrix<ACTION_LOG_STD_SPEC>& action_log_std, RNG& rng, typename DEVICE::index_t pos, typename DEVICE::index_t env_i){
+        using SPEC = typename DATASET_SPEC::SPEC;
+        using T = typename SPEC::TYPE_POLICY::DEFAULT;
+        using TI = typename SPEC::TI;
+        constexpr TI N_AGENTS = SPEC::ENVIRONMENT::N_AGENTS;
+        constexpr TI ACTION_DIM = SPEC::ENVIRONMENT::ACTION_DIM;
+        static_assert(ACTION_DIM % N_AGENTS == 0);
+        constexpr TI PER_AGENT_ACTION_DIM = SPEC::ENVIRONMENT::ACTION_DIM/N_AGENTS;
+
+        T action_log_prob = 0;
+        for(TI action_i = 0; action_i < SPEC::ENVIRONMENT::ACTION_DIM; action_i++) {
+            T action_mean = get(actions_mean, env_i, action_i);
+
+            static_assert(ACTION_DIM == ACTION_LOG_STD_SPEC::COLS * N_AGENTS);
+            static_assert(ACTION_LOG_STD_SPEC::ROWS == 1);
+            T current_action_log_std = get(action_log_std, 0, action_i % PER_AGENT_ACTION_DIM);
+            T action_std = math::exp(device.math, current_action_log_std);
+            T action_noisy = random::normal_distribution::sample(device.random, action_mean, action_std, rng);
+            action_log_prob += random::normal_distribution::log_prob(device.random, action_mean, current_action_log_std, action_noisy);
+            set(actions, env_i, action_i, action_noisy);
+        }
+        set(dataset.action_log_probs, pos, 0, action_log_prob);
+        auto& env = get(runner.environments, 0, env_i);
+        auto& state = get(runner.states, 0, env_i);
+        auto& parameters = get(runner.env_parameters, 0, env_i);
+        typename SPEC::ENVIRONMENT::State next_state;
+        auto action = row(device, actions, env_i);
+        step(device, env, parameters, state, action, next_state, rng);
+        bool terminated_flag = terminated(device, env, parameters, next_state, rng);
+        set(dataset.terminated, pos, 0, terminated_flag);
+        T reward_value = reward(device, env, parameters, state, action, next_state, rng);
+#if !defined(__CUDA_ARCH__)
+        if(env_i == 0){
+            log_reward(device, env, parameters, state, action, next_state, rng, 33);
+        }
+#endif
+        increment(runner.episode_return, 0, env_i, reward_value);
+        set(dataset.rewards, pos, 0, reward_value);
+        increment(runner.episode_step, 0, env_i, 1);
+        bool truncated = terminated_flag || (SPEC::STEP_LIMIT > 0 && get(runner.episode_step, 0, env_i) >= SPEC::STEP_LIMIT);
+        set(dataset.truncated, pos, 0, truncated);
+        TI pos_reset = pos + SPEC::N_ENVIRONMENTS; // setting dataset.reset (truncation delayed by one step for the reset of stateful actors and critics)
+        set(dataset.all_reset, pos_reset, 0, truncated);
+        set(runner.truncated, 0, env_i, truncated);
+        state = next_state;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/components/on_policy_runner/persist.h

@@ -0,0 +1,61 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_COMPONENTS_ON_POLICY_RUNNER_PERSIST_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_COMPONENTS_ON_POLICY_RUNNER_PERSIST_H
+#include "on_policy_runner.h"
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template <typename DEVICE, typename SPEC, typename GROUP>
+    void save(DEVICE& device, rl::components::on_policy_runner::Dataset<SPEC>& dataset, GROUP& group){
+        save(device, dataset.all_observations, group, "all_observations");
+        save(device, dataset.all_observations_privileged, group, "all_observations_privileged");
+        save(device, dataset.scalar_data, group, "data");
+    }
+    template <typename DEVICE, typename SPEC, typename GROUP>
+    bool load(DEVICE& device, rl::components::on_policy_runner::Dataset<SPEC>& dataset, GROUP& group){
+        bool success = load(device, dataset.all_observations, group, "all_observations");
+        success &= load(device, dataset.all_observations_privileged, group, "all_observations_privileged");
+        success &= load(device, dataset.scalar_data, group, "data");
+        return success;
+    }
+    template <typename DEVICE, typename SPEC, typename GROUP>
+    void save(DEVICE& device, rl::components::OnPolicyRunner<SPEC>& runner, GROUP& group){
+        using TI = typename DEVICE::index_t;
+        Tensor<tensor::Specification<TI, TI, tensor::Shape<TI, 1>>> step_tensor;
+        malloc(device, step_tensor);
+        set(device, step_tensor, runner.step, 0);
+        save(device, step_tensor, group, "step");
+        free(device, step_tensor);
+        auto policy_state_group = create_group(device, group, "policy_state");
+        save(device, runner.policy_state, policy_state_group);
+        save_binary(device, &get(runner.states, 0, 0), SPEC::N_ENVIRONMENTS, group, "states");
+        save_binary(device, &get(runner.env_parameters, 0, 0), SPEC::N_ENVIRONMENTS, group, "env_parameters");
+        save_binary(device, &get(runner.environments, 0, 0), SPEC::N_ENVIRONMENTS, group, "environments");
+        save(device, runner.truncated, group, "truncated");
+        save(device, runner.episode_step, group, "episode_step");
+        save(device, runner.episode_return, group, "episode_return");
+    }
+    template <typename DEVICE, typename SPEC, typename GROUP>
+    bool load(DEVICE& device, rl::components::OnPolicyRunner<SPEC>& runner, GROUP& group){
+        using TI = typename DEVICE::index_t;
+        Tensor<tensor::Specification<TI, TI, tensor::Shape<TI, 1>>> step_tensor;
+        malloc(device, step_tensor);
+        bool success = load(device, step_tensor, group, "step");
+        runner.step = get(device, step_tensor, 0);
+        free(device, step_tensor);
+        auto policy_state_group = get_group(device, group, "policy_state");
+        success &= load(device, runner.policy_state, policy_state_group);
+        success &= load_binary(device, &get(runner.states, 0, 0), SPEC::N_ENVIRONMENTS, group, "states");
+        success &= load_binary(device, &get(runner.env_parameters, 0, 0), SPEC::N_ENVIRONMENTS, group, "env_parameters");
+        success &= load_binary(device, &get(runner.environments, 0, 0), SPEC::N_ENVIRONMENTS, group, "environments");
+        success &= load(device, runner.truncated, group, "truncated");
+        success &= load(device, runner.episode_step, group, "episode_step");
+        success &= load(device, runner.episode_return, group, "episode_return");
+#ifdef RL_TOOLS_DEBUG_RL_COMPONENTS_ON_POLICY_RUNNER_CHECK_INIT
+        runner.initialized = true;
+#endif
+        return success;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/components/operations_cpu.h

@@ -0,0 +1,2 @@
+#include "off_policy_runner/operations_cpu.h"
+#include "replay_buffer/operations_cpu.h"

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/components/operations_cpu_accelerate.h

@@ -0,0 +1 @@
+#include "off_policy_runner/operations_cpu_accelerate.h"

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/components/operations_cpu_mkl.h

@@ -0,0 +1 @@
+#include "off_policy_runner/operations_cpu_mkl.h"

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

@@ -0,0 +1,2 @@
+#include "replay_buffer/operations_generic.h"
+#include "off_policy_runner/operations_generic.h"

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/components/replay_buffer/operations_cpu.h

@@ -0,0 +1,9 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_COMPONENTS_REPLAY_BUFFER_OPERATIONS_CPU_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_COMPONENTS_REPLAY_BUFFER_OPERATIONS_CPU_H
+
+#include "replay_buffer.h"
+#include "operations_generic.h"
+
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/components/replay_buffer/operations_generic.h

@@ -0,0 +1,123 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_COMPONENTS_REPLAY_BUFFER_OPERATIONS_GENERIC_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_COMPONENTS_REPLAY_BUFFER_OPERATIONS_GENERIC_H
+
+#include "replay_buffer.h"
+#include "../../../utils/generic/memcpy.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools {
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void update_views(DEVICE& device, rl::components::ReplayBuffer<SPEC>& rb){
+        typename DEVICE::index_t offset = 0;
+        rb.observations                 = view(device, rb.data, matrix::ViewSpec<SPEC::CAPACITY, SPEC::OBSERVATION_DIM           >{}, 0, offset); offset += SPEC::ASYMMETRIC_OBSERVATIONS ? SPEC::OBSERVATION_DIM : 0;
+        rb.observations_privileged      = view(device, rb.data, matrix::ViewSpec<SPEC::CAPACITY, SPEC::OBSERVATION_DIM_PRIVILEGED>{}, 0, offset); offset += SPEC::OBSERVATION_DIM_PRIVILEGED;
+        rb.actions                      = view(device, rb.data, matrix::ViewSpec<SPEC::CAPACITY, SPEC::ACTION_DIM                >{}, 0, offset); offset += SPEC::ACTION_DIM;
+        rb.rewards                      = view(device, rb.data, matrix::ViewSpec<SPEC::CAPACITY, 1                               >{}, 0, offset); offset += 1;
+        rb.next_observations            = view(device, rb.data, matrix::ViewSpec<SPEC::CAPACITY, SPEC::OBSERVATION_DIM           >{}, 0, offset); offset += SPEC::ASYMMETRIC_OBSERVATIONS ? SPEC::OBSERVATION_DIM : 0;
+        rb.next_observations_privileged = view(device, rb.data, matrix::ViewSpec<SPEC::CAPACITY, SPEC::OBSERVATION_DIM_PRIVILEGED>{}, 0, offset); offset += SPEC::OBSERVATION_DIM_PRIVILEGED;
+        rb.terminated                   = view(device, rb.data, matrix::ViewSpec<SPEC::CAPACITY, 1                               >{}, 0, offset); offset += 1;
+        rb.truncated                    = view(device, rb.data, matrix::ViewSpec<SPEC::CAPACITY, 1                               >{}, 0, offset);
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, rl::components::ReplayBuffer<SPEC>& rb) {
+        malloc(device, rb.data);
+        malloc(device, rb.episode_start);
+        update_views(device, rb);
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, rl::components::ReplayBuffer<SPEC>& rb) {
+        free(device, rb.data);
+        free(device, rb.episode_start);
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, rl::components::ReplayBufferWithStates<SPEC>& rb) {
+        malloc(device, (rl::components::ReplayBuffer<typename SPEC::BASE_SPEC>&)rb);
+        malloc(device, rb.states);
+        malloc(device, rb.next_states);
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, rl::components::ReplayBufferWithStates<SPEC>& rb) {
+        free(device, (rl::components::ReplayBuffer<typename SPEC::BASE_SPEC>&)rb);
+        free(device, rb.states);
+        free(device, rb.next_states);
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void init(DEVICE& device, rl::components::ReplayBuffer<SPEC>& rb) {
+        update_views(device, rb);
+        rb.full = false;
+        rb.position = 0;
+        rb.current_episode_start = 0;
+    }
+    template <typename DEVICE, typename SPEC, typename STATE, typename OBSERVATION_SPEC, typename OBSERVATION_PRIVILEGED_SPEC, typename ACTION_SPEC, typename NEXT_OBSERVATION_SPEC, typename NEXT_OBSERVATION_PRIVILEGED_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void add(DEVICE& device, rl::components::ReplayBuffer<SPEC>& buffer, const STATE& state, const Matrix<OBSERVATION_SPEC>& observation, const Matrix<OBSERVATION_PRIVILEGED_SPEC>& observation_privileged, const Matrix<ACTION_SPEC>& action, typename SPEC::TYPE_POLICY::DEFAULT reward, const STATE& next_state, const Matrix<NEXT_OBSERVATION_SPEC>& next_observation, const Matrix<NEXT_OBSERVATION_PRIVILEGED_SPEC>& next_observation_privileged, const bool terminated, const bool truncated) {
+        // todo: change to memcpy?
+        for(typename DEVICE::index_t i = 0; i < SPEC::OBSERVATION_DIM; i++) {
+            set(buffer.observations, buffer.position, i, get(observation, 0, i));
+            set(buffer.next_observations, buffer.position, i, get(next_observation, 0, i));
+        }
+        for(typename DEVICE::index_t i = 0; i < SPEC::OBSERVATION_DIM_PRIVILEGED; i++) {
+            set(buffer.observations_privileged, buffer.position, i, get(observation_privileged, 0, i));
+            set(buffer.next_observations_privileged, buffer.position, i, get(next_observation_privileged, 0, i));
+        }
+        for(typename DEVICE::index_t i = 0; i < SPEC::ACTION_DIM; i++) {
+            set(buffer.actions, buffer.position, i, get(action, 0, i));
+        }
+        set(buffer.rewards, buffer.position, 0, reward);
+        set(buffer.terminated, buffer.position, 0, terminated);
+        set(buffer.truncated, buffer.position, 0, truncated);
+        set(device, buffer.episode_start, buffer.current_episode_start, buffer.position);
+        buffer.position = (buffer.position + 1) % SPEC::CAPACITY;
+        if(truncated){
+            buffer.current_episode_start = buffer.position;
+        }
+        if(buffer.position == 0 && !buffer.full) {
+            buffer.full = true;
+        }
+//        add_scalar(device, device.logger, "replay_buffer/position", (typename SPEC::T)(buffer.full ? SPEC::CAPACITY : buffer.position), 1000);
+    }
+    template <typename DEVICE, typename SPEC, typename STATE, typename OBSERVATION_SPEC, typename OBSERVATION_PRIVILEGED_SPEC, typename ACTION_SPEC, typename NEXT_OBSERVATION_SPEC, typename NEXT_OBSERVATION_PRIVILEGED_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void add(DEVICE& device, rl::components::ReplayBufferWithStates<SPEC>& buffer, const STATE& state, const Matrix<OBSERVATION_SPEC>& observation, const Matrix<OBSERVATION_PRIVILEGED_SPEC>& observation_privileged, const Matrix<ACTION_SPEC>& action, const typename SPEC::T reward, const STATE& next_state, const Matrix<NEXT_OBSERVATION_SPEC>& next_observation, const Matrix<NEXT_OBSERVATION_PRIVILEGED_SPEC>& next_observation_privileged, const bool terminated, const bool truncated) {
+        set(buffer.states, buffer.position, 0, state);
+        set(buffer.next_states, buffer.position, 0, next_state);
+        add(device, (rl::components::ReplayBuffer<typename SPEC::BASE_SPEC>&) buffer, state, observation, observation_privileged, action, reward, next_state, next_observation, next_observation_privileged, terminated, truncated);
+    }
+    template <typename SOURCE_DEVICE, typename TARGET_DEVICE, typename SOURCE_SPEC, typename TARGET_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void copy(SOURCE_DEVICE& source_device, TARGET_DEVICE& target_device, rl::components::ReplayBuffer<SOURCE_SPEC>& source, rl::components::ReplayBuffer<TARGET_SPEC>& target) {
+        copy(source_device, target_device, source.data, target.data);
+        copy(source_device, target_device, source.episode_start, target.episode_start);
+        target.full = source.full;
+        target.position = source.position;
+    }
+
+    template <typename DEVICE, typename SPEC_1, typename SPEC_2>
+    RL_TOOLS_FUNCTION_PLACEMENT typename SPEC_1::TYPE_POLICY::DEFAULT abs_diff(DEVICE& device, rl::components::ReplayBuffer<SPEC_1>& b1, rl::components::ReplayBuffer<SPEC_2>& b2) {
+        typename SPEC_1::TYPE_POLICY::DEFAULT acc = 0;
+        acc += abs_diff(device, b1.observations, b2.observations);
+        acc += abs_diff(device, b1.observations_privileged, b2.observations_privileged);
+        acc += abs_diff(device, b1.actions, b2.actions);
+        acc += abs_diff(device, b1.rewards, b2.rewards);
+        acc += abs_diff(device, b1.next_observations, b2.next_observations);
+        acc += abs_diff(device, b1.next_observations_privileged, b2.next_observations_privileged);
+        acc += abs_diff(device, b1.terminated, b2.terminated);
+        acc += abs_diff(device, b1.truncated, b2.truncated);
+        return acc;
+    }
+    template <typename DEVICE, typename SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT void recalculate_rewards(DEVICE& device, rl::components::ReplayBufferWithStates<SPEC>& buffer, const typename SPEC::ENVIRONMENT& env, RNG& rng) {
+        using TI = typename DEVICE::index_t;
+        using ENVIRONMENT = typename SPEC::ENVIRONMENT;
+        using STATE = typename ENVIRONMENT::State;
+        for(TI position_i = 0; position_i < SPEC::BASE_SPEC::CAPACITY; position_i++){
+            auto prev_r = get(buffer.rewards, position_i, 0);
+            STATE& state = get(buffer.states, position_i, 0);
+            STATE& next_state = get(buffer.next_states, position_i, 0);
+            auto action = row(device, buffer.actions, position_i);
+            auto r = reward(device, env, state, action, next_state, rng);
+            set(buffer.rewards, position_i, 0, r);
+        }
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/components/replay_buffer/persist.h

@@ -0,0 +1,57 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_COMPONENTS_REPLAY_BUFFER_PERSIST_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_COMPONENTS_REPLAY_BUFFER_PERSIST_H
+
+#include "replay_buffer.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template <typename DEVICE, typename SPEC, typename GROUP>
+    void save(DEVICE& device, rl::components::ReplayBuffer<SPEC>& rb, GROUP& group) {
+        using TI = typename SPEC::TI;
+        static_assert(decltype(rb.rewards)::COLS == 1);
+        static_assert(decltype(rb.terminated)::COLS == 1);
+        static_assert(decltype(rb.truncated)::COLS == 1);
+        save(device, rb.observations, group, "observations");
+        save(device, rb.actions, group, "actions");
+        save(device, rb.rewards, group, "rewards");
+        save(device, rb.next_observations, group, "next_observations");
+        save(device, rb.terminated, group, "terminated");
+        save(device, rb.truncated, group, "truncated");
+
+        save(device, rb.episode_start, group, "episode_start");
+
+        Tensor<tensor::Specification<decltype(rb.position), TI, tensor::Shape<TI, 1>, false>> position, full;
+        set(device, position, rb.position, 0);
+        save(device, position , group, "position");
+
+        set(device, full, rb.full, 0);
+        save(device, full, group, "full");
+    }
+    template <typename DEVICE, typename SPEC, typename GROUP>
+    bool load(DEVICE& device, rl::components::ReplayBuffer<SPEC>& rb, GROUP& group) {
+        using TI = typename SPEC::TI;
+        static_assert(decltype(rb.rewards)::COLS == 1);
+        static_assert(decltype(rb.terminated)::COLS == 1);
+        static_assert(decltype(rb.truncated)::COLS == 1);
+        bool success = load(device, rb.observations, group, "observations");
+        success &= load(device, rb.actions, group, "actions");
+        success &= load(device, rb.rewards, group, "rewards");
+        success &= load(device, rb.next_observations, group, "next_observations");
+        success &= load(device, rb.terminated, group, "terminated");
+        success &= load(device, rb.truncated, group, "truncated");
+        success &= load(device, rb.episode_start, group, "episode_start");
+
+        Tensor<tensor::Specification<decltype(rb.position), TI, tensor::Shape<TI, 1>, false>> position, full;
+        success &= load(device, position, group, "position");
+        rb.position = get(device, position, 0);
+
+        success &= load(device, full, group, "full");
+        rb.full = get(device, full, 0);
+        return success;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/components/replay_buffer/replay_buffer.h

@@ -0,0 +1,93 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_COMPONENTS_REPLAY_BUFFER_REPLAY_BUFFER_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_COMPONENTS_REPLAY_BUFFER_REPLAY_BUFFER_H
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::rl::components::replay_buffer{
+    template<typename T_TYPE_POLICY, typename T_TI, T_TI T_OBSERVATION_DIM, T_TI T_OBSERVATION_DIM_PRIVILEGED, bool T_ASYMMETRIC_OBSERVATIONS, T_TI T_ACTION_DIM, T_TI T_CAPACITY, bool T_DYNAMIC_ALLOCATION=true>
+    struct Specification{
+        using TYPE_POLICY = T_TYPE_POLICY;
+        using TI = T_TI;
+        static constexpr TI OBSERVATION_DIM = T_OBSERVATION_DIM;
+        static constexpr bool ASYMMETRIC_OBSERVATIONS = T_ASYMMETRIC_OBSERVATIONS && T_OBSERVATION_DIM_PRIVILEGED > 0;
+        static_assert(ASYMMETRIC_OBSERVATIONS == T_ASYMMETRIC_OBSERVATIONS, "ASYMMETRIC_OBSERVATIONS set but not available in the environment");
+        static constexpr TI OBSERVATION_DIM_PRIVILEGED = ASYMMETRIC_OBSERVATIONS ? T_OBSERVATION_DIM_PRIVILEGED : OBSERVATION_DIM;
+        static constexpr TI OBSERVATION_DIM_PRIVILEGED_ACTUAL = ASYMMETRIC_OBSERVATIONS ? T_OBSERVATION_DIM_PRIVILEGED : 0;
+        static constexpr TI ACTION_DIM = T_ACTION_DIM;
+        static constexpr TI CAPACITY = T_CAPACITY;
+        static constexpr bool DYNAMIC_ALLOCATION = T_DYNAMIC_ALLOCATION;
+    };
+
+    template<typename T_ENVIRONMENT, typename T_BASE_SPEC>
+    struct SpecificationWithStates: T_BASE_SPEC{
+        using BASE_SPEC = T_BASE_SPEC;
+        using ENVIRONMENT = T_ENVIRONMENT;
+    };
+
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::rl::components {
+    template <typename T_SPEC>
+    struct ReplayBuffer {
+        using SPEC = T_SPEC;
+        using TYPE_POLICY = typename SPEC::TYPE_POLICY;
+        using T = typename TYPE_POLICY::DEFAULT;
+        using TI = typename SPEC::TI;
+        static constexpr TI CAPACITY = SPEC::CAPACITY;
+        static constexpr TI DATA_COLS = SPEC::OBSERVATION_DIM + SPEC::OBSERVATION_DIM_PRIVILEGED_ACTUAL + SPEC::ACTION_DIM + 1 + SPEC::OBSERVATION_DIM + SPEC::OBSERVATION_DIM_PRIVILEGED_ACTUAL + 1 + 1;
+
+        // mem
+        Matrix<matrix::Specification<T, TI, SPEC::CAPACITY, DATA_COLS, SPEC::DYNAMIC_ALLOCATION>> data;
+
+        TI current_episode_start = 0;
+        TI position = 0;
+        bool full = false;
+
+        // views
+        template<typename SPEC::TI DIM>
+        using DATA_VIEW = typename decltype(data)::template VIEW<CAPACITY, DIM>;
+
+        DATA_VIEW<SPEC::OBSERVATION_DIM> observations;
+        DATA_VIEW<SPEC::OBSERVATION_DIM_PRIVILEGED> observations_privileged;
+        DATA_VIEW<SPEC::ACTION_DIM> actions;
+        DATA_VIEW<1> rewards;
+        DATA_VIEW<SPEC::OBSERVATION_DIM> next_observations;
+        DATA_VIEW<SPEC::OBSERVATION_DIM_PRIVILEGED> next_observations_privileged;
+        DATA_VIEW<1> terminated;
+        DATA_VIEW<1> truncated;
+
+        Tensor<tensor::Specification<TI, TI, tensor::Shape<TI, SPEC::CAPACITY>, SPEC::DYNAMIC_ALLOCATION>> episode_start;
+    };
+
+    template <typename T_SPEC>
+    struct ReplayBufferWithStates: ReplayBuffer<typename T_SPEC::BASE_SPEC> {
+        // mem
+        using SPEC = typename T_SPEC::BASE_SPEC;
+        using ENVIRONMENT = typename T_SPEC::ENVIRONMENT;
+        using TYPE_POLICY = typename SPEC::TYPE_POLICY;
+        using T = typename TYPE_POLICY::DEFAULT;
+        using TI = typename SPEC::TI;
+        Matrix<matrix::Specification<typename ENVIRONMENT::State, TI, SPEC::CAPACITY, 1, SPEC::DYNAMIC_ALLOCATION>> states;
+        Matrix<matrix::Specification<typename ENVIRONMENT::State, TI, SPEC::CAPACITY, 1, SPEC::DYNAMIC_ALLOCATION>> next_states;
+    };
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::rl::components::replay_buffer{
+    template <typename T_SPEC, auto T_SIZE>
+    struct SetSpecification{
+        static constexpr auto SIZE = T_SIZE;
+    };
+    template <typename T_SET_SPEC>
+    struct Set{
+        using SET_SPEC = T_SET_SPEC;
+        using SPEC = typename SET_SPEC::SPEC;
+        ReplayBuffer<SPEC> replay_buffers[SET_SPEC::SIZE];
+    };
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/components/running_normalizer/operations_generic.h

@@ -0,0 +1,76 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_COMPONENTS_RUNNING_NORMALIZER_OPERATIONS_GENERIC_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_COMPONENTS_RUNNING_NORMALIZER_OPERATIONS_GENERIC_H
+
+#include "running_normalizer.h"
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, rl::components::RunningNormalizer<SPEC>& normalizer){
+        malloc(device, normalizer.mean);
+        malloc(device, normalizer.std);
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, rl::components::RunningNormalizer<SPEC>& normalizer){
+        free(device, normalizer.mean);
+        free(device, normalizer.std);
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void init(DEVICE& device, rl::components::RunningNormalizer<SPEC>& normalizer){
+        normalizer.age = 0;
+        set_all(device, normalizer.mean, 0);
+        set_all(device, normalizer.std, 1);
+    }
+    template <typename DEVICE, typename SPEC, typename DATA_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void update(DEVICE& device, rl::components::RunningNormalizer<SPEC>& normalizer, Matrix<DATA_SPEC>& data){
+        // Note: data should have >> 2 rows; subsequent calls should have the same number of rows to not skeew the mean/std
+        // todo: take advantage of the data coming in batches
+        static_assert(DATA_SPEC::COLS == SPEC::DIM, "Data dimension must match normalizer dimension");
+        using T = typename SPEC::TYPE_POLICY::DEFAULT;
+        using TI = typename DEVICE::index_t;
+        constexpr TI DATA_SIZE = DATA_SPEC::ROWS;
+        static_assert(DATA_SIZE > 1, "Data size must be greater than 1 and should be much greated than one");
+        normalizer.age++;
+        for(TI col_i = 0; col_i < SPEC::DIM; col_i++){
+            auto column = col(device, data, col_i);
+            T data_mean = mean(device, column);
+            T data_std = std(device, column);
+            T mean = get(normalizer.mean, 0, col_i);
+            T new_mean = mean + (data_mean - mean)/(normalizer.age);
+            T std = get(normalizer.std, 0, col_i);
+            T new_std = std + (data_std - std)/(normalizer.age);
+            set(normalizer.mean, 0, col_i, new_mean);
+            set(normalizer.std, 0, col_i, new_std);
+        }
+    }
+    template <typename DEVICE, typename SPEC, typename DATA_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void normalize(DEVICE& device, rl::components::RunningNormalizer<SPEC>& normalizer, Matrix<DATA_SPEC>& data){
+        static_assert(DATA_SPEC::COLS == SPEC::DIM, "Data dimension must match normalizer dimension");
+        normalize(device, normalizer.mean, normalizer.std, data, data);
+    }
+    template <typename DEVICE, typename SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void normalize(DEVICE& device, rl::components::RunningNormalizer<SPEC>& normalizer, Matrix<INPUT_SPEC>& input, Matrix<OUTPUT_SPEC>& output){
+        static_assert(containers::check_structure<INPUT_SPEC, OUTPUT_SPEC>);
+        static_assert(INPUT_SPEC::COLS == SPEC::DIM, "Data dimension must match normalizer dimension");
+        normalize(device, normalizer.mean, normalizer.std, input, output);
+    }
+    template <typename SOURCE_DEVICE, typename TARGET_DEVICE, typename SOURCE_SPEC, typename TARGET_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void copy(SOURCE_DEVICE& source_device, TARGET_DEVICE& target_device, rl::components::RunningNormalizer<SOURCE_SPEC>& source, rl::components::RunningNormalizer<TARGET_SPEC>& target){
+        static_assert(SOURCE_SPEC::DIM == TARGET_SPEC::DIM, "copy: source and target normalizers must have the same dimension");
+        copy(source_device, target_device, source.mean, target.mean);
+        copy(source_device, target_device, source.std, target.std);
+        target.age = source.age;
+    }
+    template <typename DEVICE, typename SPEC_1, typename SPEC_2>
+    RL_TOOLS_FUNCTION_PLACEMENT typename SPEC_1::TYPE_POLICY::DEFAULT abs_diff(DEVICE& device, rl::components::RunningNormalizer<SPEC_1>& n1, rl::components::RunningNormalizer<SPEC_2>& n2){
+        using T = typename SPEC_1::TYPE_POLICY::DEFAULT;
+        T acc = 0;
+        acc += abs_diff(device, n1.mean, n2.mean);
+        acc += abs_diff(device, n1.std, n2.std);
+        acc += math::abs(device.math, (T)n1.age - (T)n2.age);
+        return acc;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif