hyperdrone 0.1.0__py3-none-any.whl

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

@@ -0,0 +1,106 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_COMPONENTS_ON_POLICY_RUNNER_ON_POLICY_RUNNER_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_COMPONENTS_ON_POLICY_RUNNER_ON_POLICY_RUNNER_H
+
+#include "../../../utils/generic/typing.h"
+#include "../../../rl/environments/observation.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::rl::components{
+    namespace on_policy_runner{
+        template <typename T_TYPE_POLICY, typename T_TI, typename T_ENVIRONMENT, typename T_POLICY_STATE, T_TI T_N_ENVIRONMENTS = 1, T_TI T_STEP_LIMIT = 0, T_TI T_N_AGENTS_PER_ENV = 1, bool T_TRUNCATE_ON_EACH_ITERATION = false, bool T_DYNAMIC_ALLOCATION=true>
+        struct Specification{
+            using TYPE_POLICY = T_TYPE_POLICY;
+            using TI = T_TI;
+            using ENVIRONMENT = T_ENVIRONMENT;
+            using POLICY_STATE = T_POLICY_STATE;
+            static constexpr TI N_ENVIRONMENTS = T_N_ENVIRONMENTS;
+            static constexpr TI STEP_LIMIT = T_STEP_LIMIT;
+            static constexpr bool ASYMMETRIC_OBSERVATIONS = !rl_tools::utils::typing::is_same_v<typename ENVIRONMENT::Observation, typename ENVIRONMENT::ObservationPrivileged>;
+            static constexpr TI N_AGENTS_PER_ENV = T_N_AGENTS_PER_ENV; // 1 for single agent, >1 for multi-agent
+            static constexpr bool DYANMIC_ALLOCATION = T_DYNAMIC_ALLOCATION;
+            static constexpr TI EPISODE_STATS_N_ENVIRONMENTS = 1;
+            static constexpr TI EPISODE_STATS_CADENCE = 100;
+            static constexpr bool TRUNCATE_ON_EACH_ITERATION = T_TRUNCATE_ON_EACH_ITERATION;
+        };
+
+        template <typename T_SPEC, typename T_SPEC::TI T_STEPS_PER_ENV, bool T_DYNAMIC_ALLOCATION = true>
+        struct DatasetSpecification{
+            using SPEC = T_SPEC;
+            using TI = typename SPEC::TI;
+            static constexpr TI STEPS_PER_ENV = T_STEPS_PER_ENV;
+            static constexpr bool ASYMMETRIC_OBSERVATIONS = SPEC::ASYMMETRIC_OBSERVATIONS;
+            static constexpr TI STEPS_TOTAL = STEPS_PER_ENV * SPEC::N_ENVIRONMENTS;
+            static constexpr TI STEPS_TOTAL_ALL = (STEPS_PER_ENV+1) * SPEC::N_ENVIRONMENTS; // +1 for the final observation
+            static constexpr bool DYNAMIC_ALLOCATION = T_DYNAMIC_ALLOCATION;
+        };
+
+        template <typename T_DATASET_SPEC>
+        struct Dataset{
+            using DATASET_SPEC = T_DATASET_SPEC;
+            using SPEC = typename DATASET_SPEC::SPEC;
+            using TYPE_POLICY = typename SPEC::TYPE_POLICY;
+            using T = typename TYPE_POLICY::DEFAULT;
+            using TI = typename SPEC::TI;
+            static constexpr TI STEPS_PER_ENV = DATASET_SPEC::STEPS_PER_ENV;
+            static constexpr TI STEPS_TOTAL = DATASET_SPEC::STEPS_TOTAL;
+
+            using OBS_SHAPE = typename SPEC::ENVIRONMENT::Observation::SHAPE;
+            using OBS_PRIV_SHAPE = typename SPEC::ENVIRONMENT::ObservationPrivileged::SHAPE;
+
+            // Observation tensor storage (always flat to ensure matrix_view gives (N, DIM) rows)
+            using ALL_OBS_STORAGE_SHAPE = tensor::Shape<TI, DATASET_SPEC::STEPS_TOTAL_ALL, SPEC::ENVIRONMENT::Observation::DIM>;
+            using ALL_OBS_PRIV_STORAGE_SHAPE = tensor::Shape<TI, DATASET_SPEC::STEPS_TOTAL_ALL, SPEC::ENVIRONMENT::ObservationPrivileged::DIM>;
+            Tensor<tensor::Specification<T, TI, ALL_OBS_STORAGE_SHAPE, DATASET_SPEC::DYNAMIC_ALLOCATION>> all_observations;
+            Tensor<tensor::Specification<T, TI, ALL_OBS_PRIV_STORAGE_SHAPE, DATASET_SPEC::DYNAMIC_ALLOCATION>> all_observations_privileged;
+
+            // Scalar data (actions, rewards, flags, values, advantages)
+            static constexpr TI SCALAR_DATA_DIM = SPEC::ENVIRONMENT::ACTION_DIM * 2 + 8;
+            Matrix<matrix::Specification<T, TI, STEPS_TOTAL + SPEC::N_ENVIRONMENTS, SCALAR_DATA_DIM, DATASET_SPEC::DYNAMIC_ALLOCATION>> scalar_data;
+
+            template<TI VIEW_DIM, bool ALL = false>
+            using SCALAR_VIEW = typename decltype(scalar_data)::template VIEW<STEPS_TOTAL + (ALL ? SPEC::N_ENVIRONMENTS : 0), VIEW_DIM>;
+
+            SCALAR_VIEW<SPEC::ENVIRONMENT::ACTION_DIM> actions_mean;
+            SCALAR_VIEW<SPEC::ENVIRONMENT::ACTION_DIM> actions;
+            SCALAR_VIEW<1> action_log_probs;
+            SCALAR_VIEW<1> rewards;
+            SCALAR_VIEW<1> terminated;
+            SCALAR_VIEW<1> truncated;
+            SCALAR_VIEW<1, true> all_reset;
+            SCALAR_VIEW<1> reset; // = truncation delayed by one step for the reset of stateful actors and critics
+            SCALAR_VIEW<1, true> all_values;
+            SCALAR_VIEW<1> values;
+            SCALAR_VIEW<1> advantages;
+            SCALAR_VIEW<1> target_values;
+        };
+        template <typename TI, TI T_NUM_THREADS>
+        struct ExecutionHints{
+            static constexpr TI NUM_THREADS = T_NUM_THREADS;
+        };
+    }
+
+    template <typename T_SPEC>
+    struct OnPolicyRunner{
+        using SPEC = T_SPEC;
+        using TYPE_POLICY = typename SPEC::TYPE_POLICY;
+        using TI = typename SPEC::TI;
+
+        TI step = 0;
+
+        typename SPEC::POLICY_STATE policy_state;
+
+        Matrix<matrix::Specification<typename SPEC::ENVIRONMENT            , TI, 1, SPEC::N_ENVIRONMENTS, SPEC::DYANMIC_ALLOCATION>> environments;
+        Matrix<matrix::Specification<typename SPEC::ENVIRONMENT::Parameters, TI, 1, SPEC::N_ENVIRONMENTS, SPEC::DYANMIC_ALLOCATION>> env_parameters;
+        Matrix<matrix::Specification<typename SPEC::ENVIRONMENT::State     , TI, 1, SPEC::N_ENVIRONMENTS, SPEC::DYANMIC_ALLOCATION>> states;
+        Matrix<matrix::Specification<bool                                  , TI, 1, SPEC::N_ENVIRONMENTS, SPEC::DYANMIC_ALLOCATION>> truncated;
+        Matrix<matrix::Specification<TI                                    , TI, 1, SPEC::N_ENVIRONMENTS, SPEC::DYANMIC_ALLOCATION>> episode_step;
+        Matrix<matrix::Specification<typename TYPE_POLICY::DEFAULT         , TI, 1, SPEC::N_ENVIRONMENTS, SPEC::DYANMIC_ALLOCATION>> episode_return;
+#ifdef RL_TOOLS_DEBUG_RL_COMPONENTS_ON_POLICY_RUNNER_CHECK_INIT
+        bool initialized = false;
+#endif
+    };
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

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

@@ -0,0 +1,16 @@
+#include "on_policy_runner.h"
+#include "operations_generic_per_env.h"
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::rl::components::on_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::on_policy_runner::ExecutionHints<TI, NUM_THREADS> hints) {
+        return NUM_THREADS;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#ifndef RL_TOOLS_RL_COMPONENTS_ON_POLICY_RUNNER_OPERATIONS_CPU_DELAY_OPERATIONS_GENERIC_INCLUDE
+#include "operations_generic.h"
+#endif

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

@@ -0,0 +1,28 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_COMPONENTS_ON_POLICY_RUNNER_OPERATIONS_CPU_ACCELERATE_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_COMPONENTS_ON_POLICY_RUNNER_OPERATIONS_CPU_ACCELERATE_H
+
+#include "on_policy_runner.h"
+#define RL_TOOLS_RL_COMPONENTS_ON_POLICY_RUNNER_OPERATIONS_CPU_DELAY_OPERATIONS_GENERIC_INCLUDE
+#include "operations_cpu.h"
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::rl::components::on_policy_runner{
+    template <typename DEV_SPEC, typename DATASET_SPEC, typename RNG> // todo: make this not PPO but general policy with output distribution
+    void prologue(devices::CPU_ACCELERATE<DEV_SPEC>& device, rl::components::on_policy_runner::Dataset<DATASET_SPEC>& dataset, rl::components::OnPolicyRunner<typename DATASET_SPEC::SPEC>& runner, RNG& rng, typename devices::CPU<DEV_SPEC>::index_t step_i){
+        prologue((devices::CPU<DEV_SPEC>&)device, dataset, runner, rng, step_i);
+    }
+    template <typename DEV_SPEC, 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
+    void epilogue(devices::CPU_ACCELERATE<DEV_SPEC>& 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 devices::CPU<DEV_SPEC>::index_t step_i){
+        epilogue((devices::CPU<DEV_SPEC>&)device, dataset, runner, actions_mean, actions, action_log_std, rng, step_i);
+    }
+    // CPU_ACCELERATE + ArrayENGINE: resolves ambiguity with generic ArrayENGINE overload
+    template <typename DEV_SPEC, typename DATASET_SPEC, typename ACTIONS_MEAN_SPEC, typename ACTIONS_SPEC, typename ACTION_LOG_STD_SPEC, typename ARRAY_SPEC>
+    void epilogue(devices::CPU_ACCELERATE<DEV_SPEC>& 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, devices::generic::random::ArrayENGINE<ARRAY_SPEC>& rng, typename devices::CPU<DEV_SPEC>::index_t step_i){
+        epilogue((devices::CPU<DEV_SPEC>&)device, dataset, runner, actions_mean, actions, action_log_std, rng, step_i);
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#include "operations_generic.h"
+
+#endif

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

@@ -0,0 +1,15 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_COMPONENTS_ON_POLICY_RUNNER_OPERATIONS_CPU_MKL_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_COMPONENTS_ON_POLICY_RUNNER_OPERATIONS_CPU_MKL_H
+
+#include "on_policy_runner.h"
+#define RL_TOOLS_RL_COMPONENTS_ON_POLICY_RUNNER_OPERATIONS_CPU_DELAY_OPERATIONS_GENERIC_INCLUDE
+#include "operations_cpu.h"
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::rl::components::on_policy_runner{
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#include "operations_generic.h"
+
+#endif

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

@@ -0,0 +1,195 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_COMPONENTS_ON_POLICY_RUNNER_OPERATIONS_CUDA_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_COMPONENTS_ON_POLICY_RUNNER_OPERATIONS_CUDA_H
+
+#include "../../../devices/dummy.h"
+#include "operations_generic.h"
+#include "on_policy_runner.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    namespace rl::components::on_policy_runner{
+        template<typename DEVICE, typename SPEC, typename ENV_SPEC, typename PARAM_SPEC>
+        __global__
+        void init_kernel(DEVICE device, rl::components::OnPolicyRunner<SPEC> runner, Tensor<ENV_SPEC> environments, Tensor<PARAM_SPEC> parameters){
+            using TI = typename SPEC::TI;
+            TI env_i = threadIdx.x + blockIdx.x * blockDim.x;
+            if(env_i < SPEC::N_ENVIRONMENTS){
+                auto& src_env = get_ref(device, environments, env_i);
+                auto& src_params = get_ref(device, parameters, env_i);
+                set(runner.environments, 0, env_i, src_env);
+                set(runner.env_parameters, 0, env_i, src_params);
+            }
+        }
+        template<typename DEVICE, typename OBS_PRIV_SPEC, typename OBS_SPEC, typename SPEC, typename RNG>
+        __global__
+        void prologue_kernel(DEVICE device, Tensor<OBS_PRIV_SPEC> observations_privileged, Tensor<OBS_SPEC> observations, rl::components::OnPolicyRunner<SPEC> runner, RNG rng){
+            using TI = typename SPEC::TI;
+            TI env_i = threadIdx.x + blockIdx.x * blockDim.x;
+            static_assert(RNG::NUM_RNGS >= SPEC::N_ENVIRONMENTS, "Please increase the number of CUDA RNGs");
+            if(env_i < SPEC::N_ENVIRONMENTS){
+                auto& rng_state = get(rng.states, 0, env_i);
+                per_env::prologue(device, observations_privileged, observations, runner, rng_state, env_i);
+            }
+        }
+        template<typename DEV_SPEC, typename OBS_PRIV_SPEC, typename OBS_SPEC, typename SPEC, typename RNG>
+        void prologue(devices::CUDA<DEV_SPEC>& device, Tensor<OBS_PRIV_SPEC>& observations_privileged, Tensor<OBS_SPEC>& observations, rl::components::OnPolicyRunner<SPEC>& runner, RNG& rng, typename devices::CUDA<DEV_SPEC>::index_t step_i){
+            using DEVICE = devices::CUDA<DEV_SPEC>;
+            using TI = typename DEVICE::index_t;
+            constexpr TI BLOCKSIZE = 32;
+            constexpr TI N_BLOCKS = RL_TOOLS_DEVICES_CUDA_CEIL(SPEC::N_ENVIRONMENTS, BLOCKSIZE);
+            dim3 grid(N_BLOCKS);
+            dim3 block(BLOCKSIZE);
+            devices::cuda::TAG<DEVICE, true> tag_device{};
+            prologue_kernel<<<grid, block, 0, device.stream>>>(tag_device, observations_privileged, observations, runner, rng);
+            check_status(device);
+        }
+        // CUDA + ArrayENGINE: resolves ambiguity with generic ArrayENGINE overload
+        template<typename DEV_SPEC, typename OBS_PRIV_SPEC, typename OBS_SPEC, typename SPEC, typename ARRAY_SPEC>
+        void prologue(devices::CUDA<DEV_SPEC>& device, Tensor<OBS_PRIV_SPEC>& observations_privileged, Tensor<OBS_SPEC>& observations, rl::components::OnPolicyRunner<SPEC>& runner, devices::generic::random::ArrayENGINE<ARRAY_SPEC>& rng, typename devices::CUDA<DEV_SPEC>::index_t step_i){
+            using DEVICE = devices::CUDA<DEV_SPEC>;
+            using TI = typename DEVICE::index_t;
+            constexpr TI BLOCKSIZE = 32;
+            constexpr TI N_BLOCKS = RL_TOOLS_DEVICES_CUDA_CEIL(SPEC::N_ENVIRONMENTS, BLOCKSIZE);
+            dim3 grid(N_BLOCKS);
+            dim3 block(BLOCKSIZE);
+            devices::cuda::TAG<DEVICE, true> tag_device{};
+            prologue_kernel<<<grid, block, 0, device.stream>>>(tag_device, observations_privileged, observations, runner, rng);
+            check_status(device);
+        }
+        template<typename DEVICE, typename DATASET_SPEC, typename ACTIONS_MEAN_SPEC, typename ACTIONS_SPEC, typename ACTION_LOG_STD_SPEC, typename RNG>
+        __global__
+        void epilogue_kernel(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 DATASET_SPEC::SPEC::TI step_i){
+            using SPEC = typename DATASET_SPEC::SPEC;
+            using TI = typename SPEC::TI;
+            TI env_i = threadIdx.x + blockIdx.x * blockDim.x;
+            static_assert(RNG::NUM_RNGS >= SPEC::N_ENVIRONMENTS, "Please increase the number of CUDA RNGs");
+            if(env_i < SPEC::N_ENVIRONMENTS){
+                auto& rng_state = get(rng.states, 0, env_i);
+                TI pos = step_i * SPEC::N_ENVIRONMENTS + env_i;
+                per_env::epilogue(device, dataset, runner, actions_mean, actions, action_log_std, rng_state, pos, env_i);
+            }
+        }
+        template<typename DEV_SPEC, typename DATASET_SPEC, typename ACTIONS_MEAN_SPEC, typename ACTIONS_SPEC, typename ACTION_LOG_STD_SPEC, typename RNG>
+        void epilogue(devices::CUDA<DEV_SPEC>& 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 devices::CUDA<DEV_SPEC>::index_t step_i){
+            using DEVICE = devices::CUDA<DEV_SPEC>;
+            using SPEC = typename DATASET_SPEC::SPEC;
+            using TI = typename DEVICE::index_t;
+            constexpr TI BLOCKSIZE = 32;
+            constexpr TI N_BLOCKS = RL_TOOLS_DEVICES_CUDA_CEIL(SPEC::N_ENVIRONMENTS, BLOCKSIZE);
+            dim3 grid(N_BLOCKS);
+            dim3 block(BLOCKSIZE);
+            devices::cuda::TAG<DEVICE, true> tag_device{};
+            epilogue_kernel<<<grid, block, 0, device.stream>>>(tag_device, dataset, runner, actions_mean, actions, action_log_std, rng, step_i);
+            check_status(device);
+        }
+        // CUDA + ArrayENGINE: resolves ambiguity with generic ArrayENGINE overload
+        template<typename DEV_SPEC, typename DATASET_SPEC, typename ACTIONS_MEAN_SPEC, typename ACTIONS_SPEC, typename ACTION_LOG_STD_SPEC, typename ARRAY_SPEC>
+        void epilogue(devices::CUDA<DEV_SPEC>& 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, devices::generic::random::ArrayENGINE<ARRAY_SPEC>& rng, typename devices::CUDA<DEV_SPEC>::index_t step_i){
+            using DEVICE = devices::CUDA<DEV_SPEC>;
+            using SPEC = typename DATASET_SPEC::SPEC;
+            using TI = typename DEVICE::index_t;
+            constexpr TI BLOCKSIZE = 32;
+            constexpr TI N_BLOCKS = RL_TOOLS_DEVICES_CUDA_CEIL(SPEC::N_ENVIRONMENTS, BLOCKSIZE);
+            dim3 grid(N_BLOCKS);
+            dim3 block(BLOCKSIZE);
+            devices::cuda::TAG<DEVICE, true> tag_device{};
+            epilogue_kernel<<<grid, block, 0, device.stream>>>(tag_device, dataset, runner, actions_mean, actions, action_log_std, rng, step_i);
+            check_status(device);
+        }
+        template<typename DEVICE, typename DATASET_SPEC, typename SPEC, typename RNG>
+        __global__
+        void final_observation_kernel(DEVICE device, rl::components::on_policy_runner::Dataset<DATASET_SPEC> dataset, rl::components::OnPolicyRunner<SPEC> runner, RNG rng){
+            using TI = typename SPEC::TI;
+            TI env_i = threadIdx.x + blockIdx.x * blockDim.x;
+            if(env_i < SPEC::N_ENVIRONMENTS){
+                auto& rng_state = get(rng.states, 0, 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);
+                auto obs_slice = view(device, dataset.all_observations, (TI)(DATASET_SPEC::STEPS_PER_ENV * SPEC::N_ENVIRONMENTS + env_i));
+                auto obs_matrix = matrix_view(device, obs_slice);
+                observe(device, env, parameters, state, typename SPEC::ENVIRONMENT::Observation{}, obs_matrix, rng_state);
+                auto obs_priv_slice = view(device, dataset.all_observations_privileged, (TI)(DATASET_SPEC::STEPS_PER_ENV * SPEC::N_ENVIRONMENTS + env_i));
+                auto obs_priv_matrix = matrix_view(device, obs_priv_slice);
+                observe(device, env, parameters, state, typename SPEC::ENVIRONMENT::ObservationPrivileged{}, obs_priv_matrix, rng_state);
+            }
+        }
+        template<typename DEVICE, typename DATASET_SPEC, typename SPEC>
+        __global__
+        void copy_truncated_to_reset_kernel(DEVICE device, rl::components::on_policy_runner::Dataset<DATASET_SPEC> dataset, rl::components::OnPolicyRunner<SPEC> runner){
+            using TI = typename SPEC::TI;
+            TI env_i = threadIdx.x + blockIdx.x * blockDim.x;
+            if(env_i < SPEC::N_ENVIRONMENTS){
+                set(dataset.reset, env_i, 0, get(runner.truncated, 0, env_i));
+            }
+        }
+    }
+    template <typename DEV_SPEC, typename SPEC, typename ENV_SPEC, typename PARAM_SPEC, typename ACTOR, typename RNG>
+    void init(devices::CUDA<DEV_SPEC>& device, rl::components::OnPolicyRunner<SPEC>& runner, Tensor<ENV_SPEC> environments, Tensor<PARAM_SPEC> parameters, ACTOR& actor, RNG& rng){
+        using DEVICE = devices::CUDA<DEV_SPEC>;
+        using TI = typename DEVICE::index_t;
+        set_all(device, runner.episode_step, 0);
+        set_all(device, runner.episode_return, 0);
+        set_all(device, runner.truncated, true);
+        constexpr TI BLOCKSIZE = 32;
+        constexpr TI N_BLOCKS = RL_TOOLS_DEVICES_CUDA_CEIL(SPEC::N_ENVIRONMENTS, BLOCKSIZE);
+        dim3 grid(N_BLOCKS);
+        dim3 block(BLOCKSIZE);
+        devices::cuda::TAG<DEVICE, true> tag_device{};
+        rl::components::on_policy_runner::init_kernel<<<grid, block, 0, device.stream>>>(tag_device, runner, environments, parameters);
+        check_status(device);
+        reset(device, actor, runner.policy_state, rng);
+#ifdef RL_TOOLS_DEBUG_RL_COMPONENTS_ON_POLICY_RUNNER_CHECK_INIT
+        runner.initialized = true;
+#endif
+    }
+    template <typename DEV_SPEC, typename DATASET_SPEC, typename ACTOR, typename ACTOR_BUFFER, typename RNG>
+    void collect(devices::CUDA<DEV_SPEC>& device, rl::components::on_policy_runner::Dataset<DATASET_SPEC>& dataset, rl::components::OnPolicyRunner<typename DATASET_SPEC::SPEC>& runner, ACTOR& actor, ACTOR_BUFFER& policy_eval_buffers, RNG& rng){
+#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 DEVICE = devices::CUDA<DEV_SPEC>;
+        using SPEC = typename DATASET_SPEC::SPEC;
+        using T = typename SPEC::TYPE_POLICY::DEFAULT;
+        using TI = typename SPEC::TI;
+        constexpr TI BLOCKSIZE = 32;
+        constexpr TI N_BLOCKS = RL_TOOLS_DEVICES_CUDA_CEIL(SPEC::N_ENVIRONMENTS, BLOCKSIZE);
+        dim3 grid(N_BLOCKS);
+        dim3 block(BLOCKSIZE);
+        devices::cuda::TAG<DEVICE, true> tag_device{};
+        if constexpr(SPEC::TRUNCATE_ON_EACH_ITERATION){
+            set_all(device, runner.truncated, true);
+        }
+        rl::components::on_policy_runner::copy_truncated_to_reset_kernel<<<grid, block, 0, device.stream>>>(tag_device, dataset, runner);
+        check_status(device);
+        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_privileged = view_range(device, dataset.all_observations_privileged, step_i*SPEC::N_ENVIRONMENTS, tensor::ViewSpec<0, SPEC::N_ENVIRONMENTS>{});
+            auto observations            = view_range(device, dataset.all_observations          , step_i*SPEC::N_ENVIRONMENTS, tensor::ViewSpec<0, SPEC::N_ENVIRONMENTS>{});
+            auto truncated_view = view(device, runner.truncated);
+            Mode<mode::sequential::ResetMask<mode::Default<>, mode::sequential::ResetMaskSpecification<decltype(truncated_view)>>> mode_reset_mask;
+            mode_reset_mask.mask = truncated_view;
+            reset(device, actor, runner.policy_state, rng, mode_reset_mask);
+            rl::components::on_policy_runner::prologue(device, observations_privileged, observations, runner, rng, step_i);
+            using OBS_SHAPE = typename SPEC::ENVIRONMENT::Observation::SHAPE;
+            using EVAL_INPUT_SHAPE = tensor::Prepend<OBS_SHAPE, SPEC::N_ENVIRONMENTS>;
+            auto observations_reshaped = reshape_row_major(device, observations, EVAL_INPUT_SHAPE{});
+            auto actions_mean_tensor = to_tensor(device, actions_mean);
+            Mode<mode::Rollout<>> mode;
+            evaluate_step(device, actor, observations_reshaped, runner.policy_state, actions_mean_tensor, policy_eval_buffers, rng, mode);
+            auto& last_layer = get_last_layer(actor);
+            auto log_std = matrix_view(device, last_layer.log_std.parameters);
+            rl::components::on_policy_runner::epilogue(device, dataset, runner, actions_mean, actions, log_std, rng, step_i);
+        }
+        rl::components::on_policy_runner::final_observation_kernel<<<grid, block, 0, device.stream>>>(tag_device, dataset, runner, rng);
+        check_status(device);
+        runner.step += SPEC::N_ENVIRONMENTS * DATASET_SPEC::STEPS_PER_ENV;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#include "operations_generic.h"
+
+#endif

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

@@ -0,0 +1,197 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_COMPONENTS_ON_POLICY_RUNNER_OPERATIONS_GENERIC_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_COMPONENTS_ON_POLICY_RUNNER_OPERATIONS_GENERIC_H
+
+#include "on_policy_runner.h"
+#include "operations_generic_per_env.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, rl::components::on_policy_runner::Dataset<SPEC>& dataset){
+        malloc(device, dataset.all_observations);
+        malloc(device, dataset.all_observations_privileged);
+        malloc(device, dataset.scalar_data);
+        using SCALAR_DATA_SPEC = typename decltype(dataset.scalar_data)::SPEC;
+        using TI = typename SPEC::SPEC::TI;
+        TI pos = 0;
+        dataset.actions_mean     = view<DEVICE, SCALAR_DATA_SPEC, decltype(dataset.actions_mean    )::ROWS, decltype(dataset.actions_mean    )::COLS>(device, dataset.scalar_data, 0, pos); pos += decltype(dataset.actions_mean    )::COLS;
+        dataset.actions          = view<DEVICE, SCALAR_DATA_SPEC, decltype(dataset.actions         )::ROWS, decltype(dataset.actions         )::COLS>(device, dataset.scalar_data, 0, pos); pos += decltype(dataset.actions         )::COLS;
+        dataset.action_log_probs = view<DEVICE, SCALAR_DATA_SPEC, decltype(dataset.action_log_probs)::ROWS, decltype(dataset.action_log_probs)::COLS>(device, dataset.scalar_data, 0, pos); pos += decltype(dataset.action_log_probs)::COLS;
+        dataset.rewards          = view<DEVICE, SCALAR_DATA_SPEC, decltype(dataset.rewards         )::ROWS, decltype(dataset.rewards         )::COLS>(device, dataset.scalar_data, 0, pos); pos += decltype(dataset.rewards         )::COLS;
+        dataset.terminated       = view<DEVICE, SCALAR_DATA_SPEC, decltype(dataset.terminated      )::ROWS, decltype(dataset.terminated      )::COLS>(device, dataset.scalar_data, 0, pos); pos += decltype(dataset.terminated      )::COLS;
+        dataset.truncated        = view<DEVICE, SCALAR_DATA_SPEC, decltype(dataset.truncated       )::ROWS, decltype(dataset.truncated       )::COLS>(device, dataset.scalar_data, 0, pos); pos += decltype(dataset.truncated       )::COLS;
+        dataset.all_reset        = view<DEVICE, SCALAR_DATA_SPEC, decltype(dataset.all_reset       )::ROWS, decltype(dataset.all_reset       )::COLS>(device, dataset.scalar_data, 0, pos);
+        dataset.reset            = view<DEVICE, SCALAR_DATA_SPEC, decltype(dataset.reset           )::ROWS, decltype(dataset.reset           )::COLS>(device, dataset.scalar_data, 0, pos); pos += decltype(dataset.reset           )::COLS;
+        dataset.all_values       = view<DEVICE, SCALAR_DATA_SPEC, decltype(dataset.all_values      )::ROWS, decltype(dataset.all_values      )::COLS>(device, dataset.scalar_data, 0, pos);
+        dataset.values           = view<DEVICE, SCALAR_DATA_SPEC, decltype(dataset.values          )::ROWS, decltype(dataset.values          )::COLS>(device, dataset.scalar_data, 0, pos); pos += decltype(dataset.values          )::COLS;
+        dataset.advantages       = view<DEVICE, SCALAR_DATA_SPEC, decltype(dataset.advantages      )::ROWS, decltype(dataset.advantages      )::COLS>(device, dataset.scalar_data, 0, pos); pos += decltype(dataset.advantages      )::COLS;
+        dataset.target_values    = view<DEVICE, SCALAR_DATA_SPEC, decltype(dataset.target_values   )::ROWS, decltype(dataset.target_values   )::COLS>(device, dataset.scalar_data, 0, pos);
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, rl::components::on_policy_runner::Dataset<SPEC>& dataset){
+        free(device, dataset.all_observations);
+        free(device, dataset.all_observations_privileged);
+        free(device, dataset.scalar_data);
+        dataset.actions_mean               ._data = nullptr;
+        dataset.actions                    ._data = nullptr;
+        dataset.action_log_probs           ._data = nullptr;
+        dataset.rewards                    ._data = nullptr;
+        dataset.terminated                 ._data = nullptr;
+        dataset.truncated                  ._data = nullptr;
+        dataset.all_values                 ._data = nullptr;
+        dataset.values                     ._data = nullptr;
+        dataset.advantages                 ._data = nullptr;
+        dataset.target_values              ._data = nullptr;
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, rl::components::OnPolicyRunner<SPEC>& runner){
+        malloc(device, runner.environments);
+        malloc(device, runner.env_parameters);
+        malloc(device, runner.states);
+        malloc(device, runner.policy_state);
+        malloc(device, runner.episode_step);
+        malloc(device, runner.episode_return);
+        malloc(device, runner.truncated);
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, rl::components::OnPolicyRunner<SPEC>& runner){
+        free(device, runner.environments);
+        free(device, runner.env_parameters);
+        free(device, runner.states);
+        free(device, runner.policy_state);
+        free(device, runner.episode_step);
+        free(device, runner.episode_return);
+        free(device, runner.truncated);
+    }
+    template <typename DEVICE, typename SPEC, typename ENV_SPEC, typename PARAM_SPEC, typename ACTOR, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT void init(DEVICE& device, rl::components::OnPolicyRunner<SPEC>& runner, Tensor<ENV_SPEC> environments, Tensor<PARAM_SPEC> parameters, ACTOR& actor, RNG& rng){
+        using TI = typename SPEC::TI;
+        set_all(device, runner.episode_step, 0);
+        set_all(device, runner.episode_return, 0);
+        set_all(device, runner.truncated, true);
+        for(TI env_i=0; env_i < SPEC::N_ENVIRONMENTS; env_i++){
+            set(runner.environments, 0, env_i, get_ref(device, environments, env_i));
+            set(runner.env_parameters, 0, env_i, get_ref(device, parameters, env_i));
+        }
+        reset(device, actor, runner.policy_state, rng);
+#ifdef RL_TOOLS_DEBUG_RL_COMPONENTS_ON_POLICY_RUNNER_CHECK_INIT
+        runner.initialized = true;
+#endif
+    }
+    namespace rl::components::on_policy_runner{
+        template <typename DEVICE, typename OBS_PRIV_SPEC, typename OBS_SPEC, typename SPEC, typename ARRAY_SPEC>
+        RL_TOOLS_FUNCTION_PLACEMENT void prologue(DEVICE& device, Tensor<OBS_PRIV_SPEC>& observations_privileged, Tensor<OBS_SPEC>& observations, rl::components::OnPolicyRunner<SPEC>& runner, devices::generic::random::ArrayENGINE<ARRAY_SPEC>& rng, typename DEVICE::index_t step_i){
+            using TI = typename SPEC::TI;
+            for(TI env_i = 0; env_i < SPEC::N_ENVIRONMENTS; env_i++){
+                auto& rng_state = get(rng.states, 0, env_i);
+                per_env::prologue(device, observations_privileged, observations, runner, rng_state, env_i);
+            }
+        }
+        template <typename DEVICE, typename DATASET_SPEC, typename ACTIONS_MEAN_SPEC, typename ACTIONS_SPEC, typename ACTION_LOG_STD_SPEC, typename ARRAY_SPEC>
+        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, devices::generic::random::ArrayENGINE<ARRAY_SPEC>& rng, typename DEVICE::index_t step_i){
+            using SPEC = typename DATASET_SPEC::SPEC;
+            using TI = typename SPEC::TI;
+            for(TI env_i = 0; env_i < SPEC::N_ENVIRONMENTS; env_i++){
+                TI pos = step_i * SPEC::N_ENVIRONMENTS + env_i;
+                auto& rng_state = get(rng.states, 0, env_i);
+                per_env::epilogue(device, dataset, runner, actions_mean, actions, action_log_std, rng_state, pos, env_i);
+            }
+        }
+        template <typename DEVICE, typename DATASET_SPEC, typename ARRAY_SPEC>
+        RL_TOOLS_FUNCTION_PLACEMENT void final_observations(DEVICE& device, rl::components::on_policy_runner::Dataset<DATASET_SPEC>& dataset, rl::components::OnPolicyRunner<typename DATASET_SPEC::SPEC>& runner, devices::generic::random::ArrayENGINE<ARRAY_SPEC>& rng){
+            using SPEC = typename DATASET_SPEC::SPEC;
+            using TI = typename SPEC::TI;
+            for(TI env_i = 0; env_i < SPEC::N_ENVIRONMENTS; env_i++){
+                auto& rng_state = get(rng.states, 0, 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);
+                auto obs_slice = view(device, dataset.all_observations, (TI)(DATASET_SPEC::STEPS_PER_ENV * SPEC::N_ENVIRONMENTS + env_i));
+                auto obs_matrix = matrix_view(device, obs_slice);
+                observe(device, env, parameters, state, typename SPEC::ENVIRONMENT::Observation{}, obs_matrix, rng_state);
+                auto obs_priv_slice = view(device, dataset.all_observations_privileged, (TI)(DATASET_SPEC::STEPS_PER_ENV * SPEC::N_ENVIRONMENTS + env_i));
+                auto obs_priv_matrix = matrix_view(device, obs_priv_slice);
+                observe(device, env, parameters, state, typename SPEC::ENVIRONMENT::ObservationPrivileged{}, obs_priv_matrix, rng_state);
+            }
+        }
+    }
+    template <typename DEVICE, typename DATASET_SPEC, typename ACTOR, typename ACTOR_BUFFER, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT void collect(DEVICE& device, rl::components::on_policy_runner::Dataset<DATASET_SPEC>& dataset, rl::components::OnPolicyRunner<typename DATASET_SPEC::SPEC>& runner, ACTOR& actor, ACTOR_BUFFER& policy_eval_buffers, RNG& rng){
+#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 T = typename SPEC::TYPE_POLICY::DEFAULT;
+        using TI = typename SPEC::TI;
+        if constexpr(SPEC::TRUNCATE_ON_EACH_ITERATION){
+            set_all(device, runner.truncated, true);
+        }
+        for (TI env_i = 0; env_i < SPEC::N_ENVIRONMENTS; env_i++) {
+            set(dataset.reset, env_i, 0, get(runner.truncated, 0, env_i));
+        }
+        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_privileged = view_range(device, dataset.all_observations_privileged, step_i*SPEC::N_ENVIRONMENTS, tensor::ViewSpec<0, SPEC::N_ENVIRONMENTS>{});
+            auto observations            = view_range(device, dataset.all_observations          , step_i*SPEC::N_ENVIRONMENTS, tensor::ViewSpec<0, SPEC::N_ENVIRONMENTS>{});
+            auto truncated_view = view(device, runner.truncated);
+            Mode<mode::sequential::ResetMask<mode::Default<>, mode::sequential::ResetMaskSpecification<decltype(truncated_view)>>> mode_reset_mask;
+            mode_reset_mask.mask = truncated_view;
+            reset(device, actor, runner.policy_state, rng, mode_reset_mask); // it is important that this happens before prologue because prologue resets the truncated flags on the runner
+            rl::components::on_policy_runner::prologue(device, observations_privileged, observations, runner, rng, step_i);
+            using OBS_SHAPE = typename SPEC::ENVIRONMENT::Observation::SHAPE;
+            using EVAL_INPUT_SHAPE = tensor::Prepend<OBS_SHAPE, SPEC::N_ENVIRONMENTS>;
+            auto observations_reshaped = reshape_row_major(device, observations, EVAL_INPUT_SHAPE{});
+            auto actions_mean_tensor = to_tensor(device, actions_mean);
+            Mode<mode::Rollout<>> mode;
+            evaluate_step(device, actor, observations_reshaped, runner.policy_state, actions_mean_tensor, policy_eval_buffers, rng, mode);
+            auto& last_layer = get_last_layer(actor);
+            auto log_std = matrix_view(device, last_layer.log_std.parameters);
+            rl::components::on_policy_runner::epilogue(device, dataset, runner, actions_mean, actions, log_std, rng, step_i);
+        }
+        rl::components::on_policy_runner::final_observations(device, dataset, runner, rng);
+        runner.step += SPEC::N_ENVIRONMENTS * DATASET_SPEC::STEPS_PER_ENV;
+    }
+    template <typename DEVICE, typename SPEC_1, typename SPEC_2>
+    RL_TOOLS_FUNCTION_PLACEMENT typename SPEC_1::SPEC::TYPE_POLICY::DEFAULT abs_diff(DEVICE& device, rl::components::on_policy_runner::Dataset<SPEC_1>& d1, rl::components::on_policy_runner::Dataset<SPEC_2>& d2){
+        using T = typename SPEC_1::SPEC::TYPE_POLICY::DEFAULT;
+        T acc = 0;
+        acc += abs_diff(device, d1.all_observations, d2.all_observations);
+        acc += abs_diff(device, d1.all_observations_privileged, d2.all_observations_privileged);
+        acc += abs_diff(device, d1.actions_mean, d2.actions_mean);
+        acc += abs_diff(device, d1.actions, d2.actions);
+        acc += abs_diff(device, d1.action_log_probs, d2.action_log_probs);
+        acc += abs_diff(device, d1.rewards, d2.rewards);
+        acc += abs_diff(device, d1.terminated, d2.terminated);
+        acc += abs_diff(device, d1.truncated, d2.truncated);
+        acc += abs_diff(device, d1.all_reset, d2.all_reset);
+        acc += abs_diff(device, d1.all_values, d2.all_values);
+        acc += abs_diff(device, d1.advantages, d2.advantages);
+        acc += abs_diff(device, d1.target_values, d2.target_values);
+        return acc;
+    }
+    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::OnPolicyRunner<SPEC_1>& r1, rl::components::OnPolicyRunner<SPEC_2>& r2){
+        using T = typename SPEC_1::TYPE_POLICY::DEFAULT;
+        using TI = typename DEVICE::index_t;
+        T acc = 0;
+        acc += math::abs(device.math, (T)r1.step - (T)r2.step);
+        acc += abs_diff(device, r1.policy_state, r2.policy_state);
+        for(TI env_i = 0; env_i < SPEC_1::N_ENVIRONMENTS; env_i++){
+            TI episode_step_r1 = get(r1.episode_step, 0, env_i);
+            TI episode_step_r2 = get(r2.episode_step, 0, env_i);
+            acc += math::abs(device.math, (T)episode_step_r1 - (T)episode_step_r2);
+        }
+        acc += abs_diff(device, r1.episode_return, r2.episode_return);
+        acc += abs_diff(device, r1.truncated, r2.truncated);
+        for(TI env_i = 0; env_i < SPEC_1::N_ENVIRONMENTS; env_i++){
+            acc += abs_diff(device, get(r1.states, 0, env_i), get(r2.states, 0, env_i));
+            acc += abs_diff(device, get(r1.env_parameters, 0, env_i), get(r2.env_parameters, 0, env_i));
+        }
+        return acc;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif