hyperverse 0.0.1__py3-none-any.whl

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

@@ -0,0 +1,225 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_COMPONENTS_OFF_POLICY_RUNNER_OFF_POLICY_RUNNER_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_COMPONENTS_OFF_POLICY_RUNNER_OFF_POLICY_RUNNER_H
+
+// Please include the file containing the environments operations before including this file
+#include "../../../rl/components/replay_buffer/replay_buffer.h"
+#include "../../../utils/generic/tuple/tuple.h"
+
+
+/* requirements
+- Multiple environments
+- Batched action inference
+
+
+*/
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::rl::components::off_policy_runner {
+    template <typename TI, TI T_NUM_THREADS>
+    struct ExecutionHints{
+        static constexpr TI NUM_THREADS = T_NUM_THREADS;
+    };
+    template <typename T_TYPE_POLICY, typename T_TI>
+    struct ParametersDefault{
+        using TYPE_POLICY = T_TYPE_POLICY;
+        using TI = T_TI;
+        static constexpr TI N_ENVIRONMENTS = 1;
+        static constexpr bool ASYMMETRIC_OBSERVATIONS = false;
+        static constexpr TI REPLAY_BUFFER_CAPACITY = 10000;
+        static constexpr TI EPISODE_STEP_LIMIT = 1000;
+        static constexpr bool COLLECT_EPISODE_STATS = false;
+        static constexpr TI EPISODE_STATS_BUFFER_SIZE = 0;
+        static constexpr bool SAMPLE_PARAMETERS = true;
+    };
+    template<typename T_TYPE_POLICY, typename T_TI, typename T_ENVIRONMENT, typename T_POLICIES, typename T_PARAMETERS, bool T_DYNAMIC_ALLOCATION=true>
+    struct Specification{
+        using TYPE_POLICY = T_TYPE_POLICY;
+        using TI = T_TI;
+        using ENVIRONMENT =  T_ENVIRONMENT;
+        static constexpr TI MAX_EPISODE_LENGTH = ENVIRONMENT::EPISODE_STEP_LIMIT;
+        using POLICIES = T_POLICIES;
+        using PARAMETERS = T_PARAMETERS;
+        static constexpr bool DYNAMIC_ALLOCATION = T_DYNAMIC_ALLOCATION;
+#ifndef _MSC_VER
+        static constexpr bool DYNAMIC_ALLOCATION_REPLAY_BUFFER = false;
+        static constexpr bool DYNAMIC_ALLOCATION_EPISODE_STATS = false;
+#else
+#ifdef RL_TOOLS_BACKEND_ENABLE_CUDA
+#error "CUDA requires static allocation of the replay buffers in the OffPolicyRunner"
+#endif
+        static constexpr bool DYNAMIC_ALLOCATION_REPLAY_BUFFER = true;
+        static constexpr bool DYNAMIC_ALLOCATION_EPISODE_STATS = true;
+#endif
+        static_assert((PARAMETERS::ASYMMETRIC_OBSERVATIONS && ENVIRONMENT::ObservationPrivileged::DIM > 0) == PARAMETERS::ASYMMETRIC_OBSERVATIONS, "ASYMMETRIC_OBSERVATIONS requested but not available in the environment");
+        static constexpr TI OBSERVATION_DIM_PRIVILEGED = PARAMETERS::ASYMMETRIC_OBSERVATIONS ? ENVIRONMENT::ObservationPrivileged::DIM : ENVIRONMENT::Observation::DIM;
+        static constexpr TI OBSERVATION_DIM_PRIVILEGED_ACTUAL = PARAMETERS::ASYMMETRIC_OBSERVATIONS ? ENVIRONMENT::ObservationPrivileged::DIM : 0;
+    };
+
+    template<typename SPEC, bool DYNAMIC_ALLOCATION = SPEC::DYNAMIC_ALLOCATION>
+    struct Buffers{
+        // todo: make the buffer exploit the observation = observation_priviliged to save memory in the case of symmetric observations
+        using T_INPUT = typename SPEC::TYPE_POLICY::template GET<numeric_types::categories::Input>;
+        using TI = typename SPEC::TI;
+
+        Matrix<matrix::Specification<T_INPUT, TI, SPEC::PARAMETERS::N_ENVIRONMENTS, SPEC::ENVIRONMENT::Observation::DIM, DYNAMIC_ALLOCATION>> observations;
+        using OBSERVATIONS_PRIVILEGED_STANDALONE = Matrix<matrix::Specification<T_INPUT, TI, SPEC::PARAMETERS::N_ENVIRONMENTS, SPEC::OBSERVATION_DIM_PRIVILEGED, DYNAMIC_ALLOCATION>>;
+        using OBSERVATIONS_PRIVILEGED_VIEW = typename decltype(observations)::template VIEW<>;
+        using OBSERVATIONS_PRIVILEGED_TYPE = rl_tools::utils::typing::conditional_t<SPEC::PARAMETERS::ASYMMETRIC_OBSERVATIONS, OBSERVATIONS_PRIVILEGED_STANDALONE, OBSERVATIONS_PRIVILEGED_VIEW>;
+        OBSERVATIONS_PRIVILEGED_TYPE observations_privileged;
+        Matrix<matrix::Specification<T_INPUT, TI, SPEC::PARAMETERS::N_ENVIRONMENTS, SPEC::ENVIRONMENT::ACTION_DIM, DYNAMIC_ALLOCATION>> actions;
+        Matrix<matrix::Specification<T_INPUT, TI, SPEC::PARAMETERS::N_ENVIRONMENTS, SPEC::ENVIRONMENT::Observation::DIM, DYNAMIC_ALLOCATION>> next_observations;
+        OBSERVATIONS_PRIVILEGED_TYPE next_observations_privileged;
+    };
+
+    // the following two base types are used to warn users when they changed the sequence length but not specify custom batch sampling parameters
+    struct Dummy{};
+    struct SequentialBatchParametersDefault{};
+    template<typename TYPE_POLICY, typename TI, TI SEQUENCE_LENGTH, typename BASE=Dummy>
+    struct SequentialBatchParameters: BASE{
+        using T = typename TYPE_POLICY::DEFAULT;
+        static constexpr bool INCLUDE_FIRST_STEP_IN_TARGETS = SEQUENCE_LENGTH > 1; // This should be false for Markovian environments (that have SEQUENCE_LENGTH == 1). For non-Markovian environments (e.g. partial observable) this should be true, especially for environments, where the first step in the environment matters (like the flag environment, where crucial information is only revealed on the first step). In that case please also make sure to configure the batch sampling accordingly: check the other BATCH_SAMPLIN_* parameters.
+        static constexpr bool ALWAYS_SAMPLE_FROM_INITIAL_STATE = SEQUENCE_LENGTH > 1;
+        static constexpr bool RANDOM_SEQ_LENGTH = SEQUENCE_LENGTH > 1;
+        static constexpr bool ENABLE_NOMINAL_SEQUENCE_LENGTH_PROBABILITY = true;
+        static constexpr T NOMINAL_SEQUENCE_LENGTH_PROBABILITY = 0.5;
+    };
+
+    template<typename T_SPEC, typename T_SPEC::TI T_SEQUENCE_LENGTH, typename T_SPEC::TI T_BATCH_SIZE, typename T_PARAMETERS = SequentialBatchParameters<typename T_SPEC::TYPE_POLICY, typename T_SPEC::TI, T_SEQUENCE_LENGTH>, bool T_DYNAMIC_ALLOCATION=true>
+    struct SequentialBatchSpecification{
+        using SPEC = T_SPEC;
+        using TYPE_POLICY = typename T_SPEC::TYPE_POLICY;
+        using TI = typename SPEC::TI;
+        static constexpr TI SEQUENCE_LENGTHH = T_SEQUENCE_LENGTH;
+        static constexpr TI PADDED_SEQUENCE_LENGTH = SEQUENCE_LENGTHH + 1;
+        static constexpr TI BATCH_SIZE = T_BATCH_SIZE;
+        static constexpr bool DYNAMIC_ALLOCATION = T_DYNAMIC_ALLOCATION;
+        using PARAMETERS = T_PARAMETERS;
+    };
+
+
+    template <typename T_SPEC>
+    struct SequentialBatch{
+        using SPEC = T_SPEC;
+        using OPR_SPEC = typename T_SPEC::SPEC;
+        using TYPE_POLICY = typename OPR_SPEC::TYPE_POLICY;
+        using T_INPUT = typename TYPE_POLICY::template GET<numeric_types::categories::Input>;
+        using T_DEFAULT = typename TYPE_POLICY::DEFAULT;
+        using TI = typename OPR_SPEC::TI;
+
+        static constexpr TI BATCH_SIZE = T_SPEC::BATCH_SIZE;
+        static constexpr TI SEQUENCE_LENGTHH = T_SPEC::SEQUENCE_LENGTHH;
+        static constexpr TI PADDED_SEQUENCE_LENGTH = T_SPEC::PADDED_SEQUENCE_LENGTH;
+        static constexpr TI OBSERVATION_DIM = OPR_SPEC::ENVIRONMENT::Observation::DIM;
+        static constexpr bool ASYMMETRIC_OBSERVATIONS = OPR_SPEC::PARAMETERS::ASYMMETRIC_OBSERVATIONS;
+        static constexpr TI OBSERVATION_DIM_PRIVILEGED = OPR_SPEC::OBSERVATION_DIM_PRIVILEGED;
+        static constexpr TI ACTION_DIM = OPR_SPEC::ENVIRONMENT::ACTION_DIM;
+        static constexpr bool DYNAMIC_ALLOCATION = T_SPEC::DYNAMIC_ALLOCATION;
+
+        static constexpr TI DATA_DIM = OBSERVATION_DIM + OPR_SPEC::OBSERVATION_DIM_PRIVILEGED_ACTUAL + ACTION_DIM;
+        Tensor<tensor::Specification<T_INPUT, TI, tensor::Shape<TI, PADDED_SEQUENCE_LENGTH, BATCH_SIZE, DATA_DIM>, DYNAMIC_ALLOCATION>> observations_actions_base;
+
+        template<typename OPR_SPEC::TI DIM>
+        using OA_VIEW_BASE = typename decltype(observations_actions_base)::template VIEW_RANGE<tensor::ViewSpec<2, DIM>>;
+
+        template<typename BASE>
+        using OA_VIEW = typename BASE::template VIEW_RANGE<tensor::ViewSpec<0, SEQUENCE_LENGTHH>>;
+
+        static constexpr TI TARGET_SEQUENCE_LENGTH = T_SPEC::PARAMETERS::INCLUDE_FIRST_STEP_IN_TARGETS ? PADDED_SEQUENCE_LENGTH : SEQUENCE_LENGTHH;
+        template<typename BASE>
+        using OA_VIEW_NEXT = typename BASE::template VIEW_RANGE<tensor::ViewSpec<0, TARGET_SEQUENCE_LENGTH>>;
+
+        OA_VIEW_BASE<OBSERVATION_DIM> observations_base;
+        OA_VIEW_BASE<OPR_SPEC::OBSERVATION_DIM_PRIVILEGED> observations_privileged_base;
+        OA_VIEW_BASE<ACTION_DIM> actions_base;
+        OA_VIEW_BASE<OPR_SPEC::OBSERVATION_DIM_PRIVILEGED + ACTION_DIM> observations_and_actions_base;
+
+        OA_VIEW <decltype(observations_base)> observations_current;
+        OA_VIEW <decltype(observations_privileged_base)> observations_privileged_current;
+        OA_VIEW <decltype(actions_base)> actions_current;
+        OA_VIEW <decltype(observations_and_actions_base)> observations_and_actions_current;
+
+        OA_VIEW_NEXT<decltype(observations_base)> observations_next;
+        OA_VIEW_NEXT<decltype(observations_privileged_base)> observations_privileged_next;
+        OA_VIEW_NEXT<decltype(actions_base)> actions_next;
+
+        Tensor<tensor::Specification<T_DEFAULT, TI, tensor::Shape<TI, SEQUENCE_LENGTHH, BATCH_SIZE, 1>, DYNAMIC_ALLOCATION>> rewards;
+        Tensor<tensor::Specification<bool, TI, tensor::Shape<TI, SEQUENCE_LENGTHH, BATCH_SIZE, 1>, DYNAMIC_ALLOCATION>> terminated;
+        Tensor<tensor::Specification<bool, TI, tensor::Shape<TI, SEQUENCE_LENGTHH, BATCH_SIZE, 1>, DYNAMIC_ALLOCATION>> reset;
+        Tensor<tensor::Specification<bool, TI, tensor::Shape<TI, PADDED_SEQUENCE_LENGTH, BATCH_SIZE, 1>, DYNAMIC_ALLOCATION>> next_reset_base;
+        typename decltype(next_reset_base)::template VIEW_RANGE<tensor::ViewSpec<0, TARGET_SEQUENCE_LENGTH>> next_reset;
+        Tensor<tensor::Specification<bool, TI, tensor::Shape<TI, SEQUENCE_LENGTHH, BATCH_SIZE, 1>, DYNAMIC_ALLOCATION>> final_step_mask;
+        Tensor<tensor::Specification<bool, TI, tensor::Shape<TI, PADDED_SEQUENCE_LENGTH, BATCH_SIZE, 1>, DYNAMIC_ALLOCATION>> next_final_step_mask_base;
+        typename decltype(next_final_step_mask_base)::template VIEW_RANGE<tensor::ViewSpec<0, TARGET_SEQUENCE_LENGTH>> next_final_step_mask;
+
+    };
+
+    template <typename T_T, typename T_TI, T_TI T_BUFFER_SIZE, bool T_DYNAMIC_ALLOCATION=true>
+    struct EpisodeStatsSpecification{
+        using T = T_T;
+        using TI = T_TI;
+        static constexpr TI BUFFER_SIZE = T_BUFFER_SIZE;
+        static constexpr bool DYNAMIC_ALLOCATION = T_DYNAMIC_ALLOCATION;
+    };
+
+    template<typename T_SPEC>
+    struct EpisodeStats{
+        using SPEC = T_SPEC;
+        using T = typename SPEC::T;
+        using TI = typename SPEC::TI;
+        Matrix<matrix::Specification<T, TI, SPEC::BUFFER_SIZE, 2, SPEC::DYNAMIC_ALLOCATION>> data;
+
+        TI next_episode_i = 0;
+        template<typename SPEC::TI DIM>
+        using STATS_VIEW = typename decltype(data)::template VIEW<SPEC::BUFFER_SIZE, DIM>;
+        STATS_VIEW<1> returns;
+        STATS_VIEW<1> steps;
+    };
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::rl::components{
+    template<typename T_SPEC>
+    struct OffPolicyRunner{
+        using SPEC = T_SPEC;
+        using TYPE_POLICY = typename SPEC::TYPE_POLICY;
+        using T = typename SPEC::TYPE_POLICY::DEFAULT;
+        using TI = typename SPEC::TI;
+        using ENVIRONMENT = typename SPEC::ENVIRONMENT;
+        using POLICIES = typename SPEC::POLICIES;
+        template <typename INPUT>
+        struct GET_STATE{
+            using CONTENT = typename INPUT::template State<SPEC::DYNAMIC_ALLOCATION>;
+        };
+        using POLICY_STATES = rl_tools::utils::MapTuple<POLICIES, GET_STATE>;
+        using REPLAY_BUFFER_SPEC = replay_buffer::Specification<TYPE_POLICY, typename SPEC::TI, SPEC::ENVIRONMENT::Observation::DIM, ENVIRONMENT::ObservationPrivileged::DIM, SPEC::PARAMETERS::ASYMMETRIC_OBSERVATIONS, SPEC::ENVIRONMENT::ACTION_DIM, SPEC::PARAMETERS::REPLAY_BUFFER_CAPACITY, SPEC::DYNAMIC_ALLOCATION_REPLAY_BUFFER>;
+        using REPLAY_BUFFER_WITH_STATES_SPEC = replay_buffer::SpecificationWithStates<ENVIRONMENT, REPLAY_BUFFER_SPEC>;
+        using REPLAY_BUFFER_TYPE = ReplayBufferWithStates<REPLAY_BUFFER_WITH_STATES_SPEC>;
+        static constexpr TI N_ENVIRONMENTS = SPEC::PARAMETERS::N_ENVIRONMENTS;
+
+        off_policy_runner::Buffers<SPEC> buffers;
+
+        TI previous_policy = 0;
+        bool previous_policy_set = false;
+
+        // todo: change to "environments"
+        Matrix<matrix::Specification<ENVIRONMENT, TI, 1, N_ENVIRONMENTS, SPEC::DYNAMIC_ALLOCATION>> envs;
+        POLICY_STATES policy_states;
+        Matrix<matrix::Specification<off_policy_runner::EpisodeStats<off_policy_runner::EpisodeStatsSpecification<T, TI, SPEC::PARAMETERS::EPISODE_STATS_BUFFER_SIZE, SPEC::DYNAMIC_ALLOCATION_EPISODE_STATS>>, TI, 1, N_ENVIRONMENTS, SPEC::DYNAMIC_ALLOCATION>> episode_stats;
+        Matrix<matrix::Specification<REPLAY_BUFFER_TYPE, TI, 1, N_ENVIRONMENTS, SPEC::DYNAMIC_ALLOCATION>> replay_buffers = {};
+
+        Matrix<matrix::Specification<typename SPEC::ENVIRONMENT::State, TI, 1, N_ENVIRONMENTS, SPEC::DYNAMIC_ALLOCATION>> states;
+        Matrix<matrix::Specification<typename SPEC::ENVIRONMENT::Parameters, TI, 1, N_ENVIRONMENTS, SPEC::DYNAMIC_ALLOCATION>> env_parameters;
+        Matrix<matrix::Specification<T, TI, 1, N_ENVIRONMENTS, SPEC::DYNAMIC_ALLOCATION>> episode_return;
+        Matrix<matrix::Specification<TI, TI, 1, N_ENVIRONMENTS, SPEC::DYNAMIC_ALLOCATION>> episode_step;
+        Matrix<matrix::Specification<bool, TI, 1, N_ENVIRONMENTS, SPEC::DYNAMIC_ALLOCATION>> truncated; // init to true !!
+#ifdef RL_TOOLS_DEBUG_RL_COMPONENTS_OFF_POLICY_RUNNER_CHECK_INIT
+        bool initialized = false;
+#endif
+    };
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif

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

@@ -0,0 +1,95 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_COMPONENTS_OFF_POLICY_RUNNER_OPERATIONS_CPU_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_COMPONENTS_OFF_POLICY_RUNNER_OPERATIONS_CPU_H
+
+#include <thread>
+
+#include "operations_generic_per_env.h"
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::rl::components::off_policy_runner{
+    constexpr auto get_num_threads(devices::ExecutionHints hints) {
+        return 1;
+    }
+    template<typename TI, TI NUM_THREADS>
+    constexpr TI get_num_threads(rl::components::off_policy_runner::ExecutionHints<TI, NUM_THREADS> hints) {
+        return NUM_THREADS;
+    }
+
+    template<typename DEV_SPEC, typename SPEC, typename RNG>
+    void prologue(devices::CPU<DEV_SPEC>& device, rl::components::OffPolicyRunner<SPEC>& runner, RNG &rng) {
+        using DEVICE = devices::CPU<DEV_SPEC>;
+        using TI = typename DEVICE::index_t;
+        std::vector<std::thread> threads;
+
+        constexpr TI NUM_THREADS = get_num_threads(typename DEVICE::EXECUTION_HINTS());
+
+        if(NUM_THREADS > 1){
+            RNG rngs[SPEC::PARAMETERS::N_ENVIRONMENTS];
+            auto base = random::uniform_int_distribution(typename DEV_SPEC::RANDOM(), 0, 1000000, rng);
+            for (TI env_i = 0; env_i < SPEC::PARAMETERS::N_ENVIRONMENTS; env_i++) {
+                init(device, rngs[env_i], base + env_i);
+            }
+
+            for (TI thread_i = 0; thread_i < NUM_THREADS; thread_i++) {
+                threads.emplace_back([NUM_THREADS, &device, thread_i, &runner, &rngs](){
+                    for (TI env_i = thread_i; env_i < SPEC::PARAMETERS::N_ENVIRONMENTS; env_i += NUM_THREADS) {
+                        prologue_per_env(device, runner, rngs[env_i], env_i);
+                    }
+                });
+            }
+
+            for (auto& thread : threads) {
+                thread.join();
+            }
+        }
+        else{
+            for (TI env_i = 0; env_i < SPEC::PARAMETERS::N_ENVIRONMENTS; env_i++) {
+                prologue_per_env(device, runner, rng, env_i);
+            }
+        }
+
+    }
+
+    template<typename DEV_SPEC, typename SPEC, typename POLICY, typename RNG>
+    void epilogue(devices::CPU<DEV_SPEC>& device, rl::components::OffPolicyRunner<SPEC>& runner, const POLICY& policy, RNG& rng){
+        using DEVICE = devices::CPU<DEV_SPEC>;
+        using TI = typename DEVICE::index_t;
+        using OPR = rl::components::OffPolicyRunner<SPEC>;
+        std::vector<std::thread> threads;
+
+        constexpr TI NUM_THREADS = get_num_threads(typename DEVICE::EXECUTION_HINTS());
+
+        if(NUM_THREADS > 1){
+            RNG rngs[OPR::N_ENVIRONMENTS];
+            auto base = random::uniform_int_distribution(typename DEV_SPEC::RANDOM(), 0, 1000000, rng);
+            for (TI env_i = 0; env_i < OPR::N_ENVIRONMENTS; env_i++) {
+                init(device, rngs[env_i], base + env_i);
+            }
+
+            for (TI thread_i = 0; thread_i < NUM_THREADS; thread_i++) {
+                threads.emplace_back([NUM_THREADS, &device, thread_i, &runner, &policy, &rngs](){
+                    for (TI env_i = thread_i; env_i < OPR::N_ENVIRONMENTS; env_i += NUM_THREADS) {
+                        epilogue_per_env(device, runner, policy, rngs[env_i], env_i);
+                    }
+                });
+            }
+
+            for (auto& thread : threads) {
+                thread.join();
+            }
+        }
+        else{
+            for (TI env_i = 0; env_i < OPR::N_ENVIRONMENTS; env_i++) {
+                epilogue_per_env(device, runner, policy, rng, env_i);
+            }
+        }
+
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#ifndef RL_TOOLS_RL_COMPONENTS_OFF_POLICY_RUNNER_OPERATIONS_CPU_DELAY_OPERATIONS_GENERIC_INCLUDE
+#include "operations_generic.h"
+#endif
+
+#endif

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

@@ -0,0 +1,21 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_COMPONENTS_OFF_POLICY_RUNNER_OPERATIONS_CPU_ACCELERATE_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_COMPONENTS_OFF_POLICY_RUNNER_OPERATIONS_CPU_ACCELERATE_H
+
+#define RL_TOOLS_RL_COMPONENTS_OFF_POLICY_RUNNER_OPERATIONS_CPU_DELAY_OPERATIONS_GENERIC_INCLUDE
+#include "operations_cpu.h"
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::rl::components::off_policy_runner{
+    template<typename DEV_SPEC, typename SPEC, typename RNG>
+    void prologue(devices::CPU_ACCELERATE<DEV_SPEC>& device, rl::components::OffPolicyRunner<SPEC>& runner, RNG &rng) {
+        prologue((devices::CPU<DEV_SPEC>&)device, runner, rng);
+    }
+    template<typename DEV_SPEC, typename SPEC, typename RNG>
+    void epilogue(devices::CPU_ACCELERATE<DEV_SPEC>& device, rl::components::OffPolicyRunner<SPEC>& runner, RNG &rng) {
+        epilogue((devices::CPU<DEV_SPEC>&)device, runner, rng);
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#include "operations_generic.h"
+#endif

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

@@ -0,0 +1,21 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_COMPONENTS_OFF_POLICY_RUNNER_OPERATIONS_CPU_MKL_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_COMPONENTS_OFF_POLICY_RUNNER_OPERATIONS_CPU_MKL_H
+
+#define RL_TOOLS_RL_COMPONENTS_OFF_POLICY_RUNNER_OPERATIONS_CPU_DELAY_OPERATIONS_GENERIC_INCLUDE
+#include "operations_cpu.h"
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::rl::components::off_policy_runner{
+    template<typename DEV_SPEC, typename SPEC, typename RNG>
+    void prologue(devices::CPU_MKL<DEV_SPEC>& device, rl::components::OffPolicyRunner<SPEC>& runner, RNG &rng) {
+        prologue((devices::CPU<DEV_SPEC>&)device, runner, rng);
+    }
+    template<typename DEV_SPEC, typename SPEC, typename POLICY, typename RNG>
+    void epilogue(devices::CPU_MKL<DEV_SPEC>& device, rl::components::OffPolicyRunner<SPEC>& runner, const POLICY& policy, RNG& rng){
+        epilogue(static_cast<devices::CPU<DEV_SPEC>&>(device), runner, policy, rng);
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#include "operations_generic.h"
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/components/off_policy_runner/operations_cuda.h

@@ -0,0 +1,237 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_COMPONENTS_OFF_POLICY_RUNNER_OPERATIONS_CUDA_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_COMPONENTS_OFF_POLICY_RUNNER_OPERATIONS_CUDA_H
+
+#include "../../../devices/dummy.h"
+#include "operations_generic.h"
+#include "off_policy_runner.h"
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    namespace rl::components::off_policy_runner{
+        template <typename DEVICE, typename SPEC>
+        __global__
+        void init_kernel(DEVICE device, rl::components::OffPolicyRunner<SPEC> runner){
+            static_assert(decltype(runner.replay_buffers)::T::SPEC::DYNAMIC_ALLOCATION == false, "The replay buffers should be statically allocated when using the OffPolicyRunner using CUDA");
+            static_assert(decltype(runner.episode_stats)::T::SPEC::DYNAMIC_ALLOCATION == false, "The episode stats should be statically allocated when using the OffPolicyRunner using CUDA");
+            using TI = typename SPEC::TI;
+            // if the episode is done (step limit activated for STEP_LIMIT > 0) or if the step is the first step for this runner, reset the environment
+            TI env_i = threadIdx.x + blockIdx.x * blockDim.x;
+            if(env_i < SPEC::PARAMETERS::N_ENVIRONMENTS){
+                auto& replay_buffer = get(runner.replay_buffers, 0, env_i);
+                replay_buffer = {}; // zero-initialize such that in debug mode the nullptr check on malloc does not trigger
+                init(device, replay_buffer);
+                auto& episode_stats = get(runner.episode_stats, 0, env_i);
+                episode_stats = {};
+                init_views(device, episode_stats);
+                init(device, episode_stats);
+                auto& env = get(runner.envs, 0, env_i);
+                env = {};
+                init(device, env);
+                auto& parameters = get(runner.env_parameters, 0, env_i);
+                initial_parameters(device, env, parameters);
+            }
+        }
+        template <bool DETERMINISTIC, typename DEVICE, typename RUNNER_SPEC, typename BATCH_SPEC, typename RNG>
+        __global__
+        void gather_batch_kernel(DEVICE device, rl::components::OffPolicyRunner<RUNNER_SPEC> runner, rl::components::off_policy_runner::SequentialBatch<BATCH_SPEC> batch, RNG rng) {
+            int dummy;
+            void* stack_ptr = &dummy;
+            // printf("Stack pointer in thread (%d, %d): %p\n", blockIdx.x, threadIdx.x, stack_ptr);
+            using TI = typename RUNNER_SPEC::TI;
+            using RUNNER = rl::components::OffPolicyRunner<RUNNER_SPEC>;
+            // if the episode is done (step limit activated for STEP_LIMIT > 0) or if the step is the first step for this runner, reset the environment
+            TI batch_step_i = threadIdx.x + blockIdx.x * blockDim.x;
+            static_assert(RNG::NUM_RNGS >= BATCH_SPEC::BATCH_SIZE, "Please increase the number of CUDA RNGs");
+            if(batch_step_i < BATCH_SPEC::BATCH_SIZE){
+                auto& rng_state = get(rng.states, 0, batch_step_i);
+                typename DEVICE::index_t env_i = DETERMINISTIC ? 0 : random::uniform_int_distribution(typename DEVICE::SPEC::RANDOM(), (TI) 0, RUNNER_SPEC::PARAMETERS::N_ENVIRONMENTS - 1, rng_state);
+                // printf("Chose env %d\n", env_i);
+                auto& replay_buffer = get(runner.replay_buffers, 0, env_i);
+                // printf("replay buffer pointer %p\n", replay_buffer.data._data);
+                // printf("replay buffer view pointer %p\n", replay_buffer.observations._data);
+                // printf("batch pointer %p\n", batch.observations_actions_next_observations._data);
+                gather_batch_step<DETERMINISTIC>(device, runner, replay_buffer, batch, batch_step_i, rng_state);
+            }
+        }
+
+        template<typename DEVICE, typename SPEC, typename RNG>
+        __global__
+        void prologue_kernel(DEVICE device, rl::components::OffPolicyRunner<SPEC> runner, RNG rng) {
+            using TI = typename SPEC::TI;
+            // if the episode is done (step limit activated for STEP_LIMIT > 0) or if the step is the first step for this runner, reset the environment
+            TI env_i = threadIdx.x + blockIdx.x * blockDim.x;
+            static_assert(RNG::NUM_RNGS >= SPEC::PARAMETERS::N_ENVIRONMENTS, "Please increase the number of CUDA RNGs");
+            if(env_i < SPEC::PARAMETERS::N_ENVIRONMENTS){
+                auto& rng_state = get(rng.states, 0, env_i);
+                prologue_per_env(device, runner, rng_state, env_i);
+            }
+        }
+        template<typename DEV_SPEC, typename SPEC, typename RNG>
+        void prologue(devices::CUDA<DEV_SPEC>& device, rl::components::OffPolicyRunner<SPEC>& runner, RNG &rng) {
+            using DEVICE = devices::CUDA<DEV_SPEC>;
+            using TI = typename SPEC::TI;
+            constexpr TI BLOCKSIZE_COLS = 32;
+            constexpr TI N_BLOCKS_COLS = RL_TOOLS_DEVICES_CUDA_CEIL(SPEC::PARAMETERS::N_ENVIRONMENTS, BLOCKSIZE_COLS);
+            dim3 grid(N_BLOCKS_COLS);
+            dim3 block(BLOCKSIZE_COLS);
+            devices::cuda::TAG<DEVICE, true> tag_device{};
+            prologue_kernel<<<grid, block, 0, device.stream>>>(tag_device, runner, rng);
+            check_status(device);
+        }
+        // template<auto POLICY_INDEX, typename DEV_SPEC, typename SPEC, typename POLICY, typename RNG>
+        // void interlude(devices::CUDA<DEV_SPEC>& device, rl::components::OffPolicyRunner<SPEC>& runner, POLICY& policy, typename POLICY::template Buffer<SPEC::PARAMETERS::N_ENVIRONMENTS>& policy_eval_buffers, RNG& rng) {
+        //     // runner struct should be on the CPU while its buffers should be on the GPU
+        //     evaluate(device, policy, runner.buffers.observations, runner.buffers.actions, policy_eval_buffers, rng);
+        // }
+
+        template<typename DEVICE, typename SPEC, typename POLICY, typename RNG>
+        __global__
+        void epilogue_kernel(DEVICE device, rl::components::OffPolicyRunner<SPEC> runner, POLICY policy, RNG rng) {
+            using TI = typename SPEC::TI;
+            using OFF_POLICY_RUNNER = rl::components::OffPolicyRunner<SPEC>;
+            static_assert(RNG::NUM_RNGS >= SPEC::PARAMETERS::N_ENVIRONMENTS, "Please increase the number of CUDA RNGs");
+            TI env_i = threadIdx.x + blockIdx.x * blockDim.x;
+            static_assert(OFF_POLICY_RUNNER::REPLAY_BUFFER_TYPE::SPEC::DYNAMIC_ALLOCATION == false, "When using the OffPolicyRunner with CUDA, the replay buffers should be statically allocated");
+            if(env_i < SPEC::PARAMETERS::N_ENVIRONMENTS){
+                POLICY dummy_policy;
+                auto& rng_state = get(rng.states, 0, env_i);
+                epilogue_per_env(device, runner, dummy_policy, rng_state, env_i);
+            }
+        }
+        template<typename DEV_SPEC, typename SPEC, typename POLICY, typename RNG>
+        void epilogue(devices::CUDA<DEV_SPEC>& device, rl::components::OffPolicyRunner<SPEC>& runner, const POLICY& policy, RNG& rng) {
+            using DEVICE = devices::CUDA<DEV_SPEC>;
+            using TI = typename SPEC::TI;
+            constexpr TI BLOCKSIZE_COLS = 32;
+            constexpr TI N_BLOCKS_COLS = RL_TOOLS_DEVICES_CUDA_CEIL(SPEC::PARAMETERS::N_ENVIRONMENTS, BLOCKSIZE_COLS);
+            dim3 grid(N_BLOCKS_COLS);
+            dim3 block(BLOCKSIZE_COLS);
+            devices::cuda::TAG<DEVICE, true> tag_device{};
+            POLICY dummy_policy; // just for type inference
+            epilogue_kernel<<<grid, block, 0, device.stream>>>(tag_device, runner, dummy_policy, rng);
+            check_status(device);
+        }
+    }
+    template<typename DEV_SPEC, typename SPEC>
+    void malloc(devices::CUDA<DEV_SPEC>& device, rl::components::OffPolicyRunner<SPEC> &runner) {
+        using DEVICE = devices::CUDA<DEV_SPEC>;
+        static_assert(!SPEC::DYNAMIC_ALLOCATION_REPLAY_BUFFER, "The replay buffer should not be dynamically allocated when using the OffPolicyRunner using CUDA");
+        static_assert(!SPEC::DYNAMIC_ALLOCATION_EPISODE_STATS, "The episode stats should not be dynamically allocated when using the OffPolicyRunner using CUDA");
+
+        malloc(device, runner.buffers);
+        malloc(device, runner.envs);
+        malloc(device, runner.states);
+        malloc(device, runner.env_parameters);
+        malloc(device, runner.episode_return);
+        malloc(device, runner.episode_step);
+        malloc(device, runner.truncated);
+        malloc(device, runner.replay_buffers);
+        malloc(device, runner.episode_stats);
+        malloc(device, runner.policy_states);
+    }
+    template<typename DEV_SPEC, typename SPEC>
+    void free(devices::CUDA<DEV_SPEC>& device, rl::components::OffPolicyRunner<SPEC> &runner) {
+        using DEVICE = devices::CUDA<DEV_SPEC>;
+        static_assert(!SPEC::DYNAMIC_ALLOCATION_REPLAY_BUFFER, "The replay buffer should not be dynamically allocated when using the OffPolicyRunner using CUDA");
+        static_assert(!SPEC::DYNAMIC_ALLOCATION_EPISODE_STATS, "The episode stats should not be dynamically allocated when using the OffPolicyRunner using CUDA");
+
+        free(device, runner.buffers);
+        free(device, runner.envs);
+        free(device, runner.states);
+        free(device, runner.env_parameters);
+        free(device, runner.episode_return);
+        free(device, runner.episode_step);
+        free(device, runner.truncated);
+        free(device, runner.replay_buffers);
+        free(device, runner.episode_stats);
+        free(device, runner.policy_states);
+    }
+    template<typename DEV_SPEC, typename SPEC>
+    void init(devices::CUDA<DEV_SPEC>& device, rl::components::OffPolicyRunner<SPEC> &runner){
+        using DEVICE = devices::CUDA<DEV_SPEC>;
+        using TI = typename DEVICE::index_t;
+        constexpr TI BLOCKSIZE_COLS = 32;
+        constexpr TI N_BLOCKS_COLS = RL_TOOLS_DEVICES_CUDA_CEIL(SPEC::PARAMETERS::N_ENVIRONMENTS, BLOCKSIZE_COLS);
+        dim3 grid(N_BLOCKS_COLS);
+        dim3 block(BLOCKSIZE_COLS);
+        devices::cuda::TAG<DEVICE, true> tag_device{};
+        rl::components::off_policy_runner::init_kernel<<<grid, block, 0, device.stream>>>(tag_device, runner);
+        check_status(device);
+
+        truncate_all(device, runner);
+        runner.previous_policy_set = false;
+        runner.previous_policy = 0;
+#ifdef RL_TOOLS_DEBUG_RL_COMPONENTS_OFF_POLICY_RUNNER_CHECK_INIT
+        runner.initialized = true;
+#endif
+    }
+    template <typename DEV_SPEC, typename SPEC, typename BATCH_SPEC, typename RNG, bool DETERMINISTIC=false>
+    void gather_batch(devices::CUDA<DEV_SPEC>& device, rl::components::OffPolicyRunner<SPEC>& runner, rl::components::off_policy_runner::SequentialBatch<BATCH_SPEC>& batch, RNG& rng) {
+        using DEVICE = devices::CUDA<DEV_SPEC>;
+        static_assert(utils::typing::is_same_v<SPEC, typename BATCH_SPEC::SPEC>);
+        using TI = typename SPEC::TI;
+        constexpr typename DEVICE::index_t BATCH_SIZE = BATCH_SPEC::BATCH_SIZE;
+        constexpr TI BLOCKSIZE_COLS = 32;
+        constexpr TI N_BLOCKS_COLS = RL_TOOLS_DEVICES_CUDA_CEIL(BATCH_SIZE, BLOCKSIZE_COLS);
+        dim3 bias_grid(N_BLOCKS_COLS);
+        dim3 bias_block(BLOCKSIZE_COLS);
+        devices::cuda::TAG<DEVICE, true> tag_device{};
+        rl::components::off_policy_runner::gather_batch_kernel<DETERMINISTIC><<<bias_grid, bias_block, 0, device.stream>>>(tag_device, runner, batch, rng);
+        check_status(device);
+    }
+    template<auto POLICY_INDEX, typename DEV_SPEC, typename SPEC, typename POLICY, typename POLICY_BUFFERS, typename RNG>
+    void step(devices::CUDA<DEV_SPEC>& device, rl::components::OffPolicyRunner<SPEC>& runner, POLICY& policy, POLICY_BUFFERS& policy_eval_buffers, RNG &rng){
+        using DEVICE = devices::CUDA<DEV_SPEC>;
+#ifdef RL_TOOLS_DEBUG_RL_COMPONENTS_OFF_POLICY_RUNNER_CHECK_INIT
+        utils::assert_exit(device, runner.initialized, "OffPolicyRunner not initialized");
+#endif
+        // static_assert(POLICY::INPUT_DIM == SPEC::ENVIRONMENT::Observation::DIM, "The policy's input dimension must match the environment's observation dimension.");
+        // static_assert(POLICY::OUTPUT_DIM == SPEC::ENVIRONMENT::ACTION_DIM, "The policy's output dimension must match the environment's action dimension.");
+        // todo: increase efficiency by removing the double observation of each state
+
+        bool policy_switch = runner.previous_policy_set && (POLICY_INDEX != runner.previous_policy);
+        if(policy_switch){
+            truncate_all(device, runner);
+        }
+
+        rl::components::off_policy_runner::prologue(device, runner, rng);
+        rl::components::off_policy_runner::interlude<POLICY_INDEX>(device, runner, policy, policy_eval_buffers, rng);
+        rl::components::off_policy_runner::epilogue(device, runner, policy, rng);
+
+        runner.previous_policy = POLICY_INDEX;
+        runner.previous_policy_set = true;
+    }
+    namespace rl::components::off_policy_runner::kernels {
+        template <typename DEV_SPEC, typename SPEC>
+        __global__
+        void update_views(devices::CUDA<DEV_SPEC> device, Matrix<SPEC> replay_buffers) {
+            using DEVICE = devices::CUDA<DEV_SPEC>;
+            static_assert(SPEC::ROWS == 1);
+            using TI = typename DEVICE::index_t;
+            TI thread_i = threadIdx.x + blockIdx.x * blockDim.x;
+            if(thread_i < SPEC::COLS) {
+                auto& rb = get(replay_buffers, 0, thread_i);
+                rl_tools::update_views(device, rb);
+            }
+        }
+    }
+    template <typename DEV_SPEC, typename SPEC>
+    void update_views(devices::CUDA<DEV_SPEC>& device, Matrix<SPEC>& replay_buffers) {
+        using DEVICE = devices::CUDA<DEV_SPEC>;
+        static_assert(SPEC::ROWS == 1);
+        using TI = typename DEVICE::index_t;
+        constexpr TI BLOCKSIZE_COLS = 32;
+        constexpr TI N_BLOCKS_COLS = RL_TOOLS_DEVICES_CUDA_CEIL(SPEC::COLS, BLOCKSIZE_COLS);
+        dim3 bias_grid(N_BLOCKS_COLS);
+        dim3 bias_block(BLOCKSIZE_COLS);
+        devices::cuda::TAG<DEVICE, true> tag_device{};
+        rl::components::off_policy_runner::kernels::update_views<<<bias_grid, bias_block, 0, device.stream>>>(tag_device, replay_buffers);
+        check_status(device);
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#include "operations_generic.h"
+
+
+#endif