hyperdrone 0.1.0__py3-none-any.whl

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

@@ -0,0 +1,193 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ALGORITHMS_PPO_OPERATIONS_CUDA_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ALGORITHMS_PPO_OPERATIONS_CUDA_H
+
+#include "ppo.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    namespace rl::algorithms::ppo::cuda{
+        template <typename DEVICE, typename DATASET_SPEC, typename PPO_PARAMETERS>
+        __global__
+        void estimate_generalized_advantages_kernel(DEVICE device, rl::components::on_policy_runner::Dataset<DATASET_SPEC> dataset, PPO_PARAMETERS){
+            using OPR_SPEC = typename DATASET_SPEC::SPEC;
+            using T = typename OPR_SPEC::TYPE_POLICY::DEFAULT;
+            using TI = typename DEVICE::index_t;
+            constexpr TI STEPS_PER_ENV = DATASET_SPEC::STEPS_PER_ENV;
+            TI env_i = threadIdx.x + blockIdx.x * blockDim.x;
+            if(env_i < OPR_SPEC::N_ENVIRONMENTS){
+                T previous_value = get(dataset.all_values, STEPS_PER_ENV * OPR_SPEC::N_ENVIRONMENTS + env_i, 0);
+                T previous_advantage = 0;
+                for(TI step_forward_i = 0; step_forward_i < STEPS_PER_ENV; step_forward_i++){
+                    TI step_backward_i = (STEPS_PER_ENV - 1 - step_forward_i);
+                    TI pos = step_backward_i * OPR_SPEC::N_ENVIRONMENTS + env_i;
+                    bool terminated_flag = get(dataset.terminated, pos, 0);
+                    bool truncated_flag = get(dataset.truncated, pos, 0);
+                    T current_step_value = get(dataset.values, pos, 0);
+                    bool terminated_actual = terminated_flag && !PPO_PARAMETERS::IGNORE_TERMINATION;
+                    T next_step_value = terminated_actual ? 0 : previous_value;
+                    T td_error = get(dataset.rewards, pos, 0) + PPO_PARAMETERS::GAMMA * next_step_value - current_step_value;
+                    if(truncated_flag){
+                        if(!terminated_flag){
+                            td_error = 0;
+                        }
+                        previous_advantage = 0;
+                    }
+                    T advantage = PPO_PARAMETERS::LAMBDA * PPO_PARAMETERS::GAMMA * previous_advantage + td_error;
+                    set(dataset.advantages, pos, 0, advantage);
+                    set(dataset.target_values, pos, 0, advantage + current_step_value);
+                    previous_advantage = advantage;
+                    previous_value = current_step_value;
+                }
+            }
+        }
+        template <typename DEVICE, typename PPO_SPEC, typename DATASET_SPEC, typename BUFFERS_SPEC, typename BATCH_ADVANTAGES_SPEC, typename BATCH_ACTIONS_SPEC, typename BATCH_ACTION_LOG_PROBS_SPEC, typename LOG_STD_SPEC>
+        __global__
+        void ppo_per_sample_kernel(DEVICE device, rl::algorithms::ppo::Buffers<BUFFERS_SPEC> ppo_buffers, Matrix<BATCH_ADVANTAGES_SPEC> batch_advantages, Matrix<BATCH_ACTIONS_SPEC> batch_actions, Matrix<BATCH_ACTION_LOG_PROBS_SPEC> batch_action_log_probs, Matrix<LOG_STD_SPEC> log_std_params){
+            using T = typename PPO_SPEC::TYPE_POLICY::DEFAULT;
+            using TI = typename DEVICE::index_t;
+            constexpr TI BATCH_SIZE = PPO_SPEC::PARAMETERS::BATCH_SIZE;
+            constexpr TI ACTION_DIM = PPO_SPEC::ENVIRONMENT::ACTION_DIM;
+            constexpr TI N_AGENTS = PPO_SPEC::ENVIRONMENT::N_AGENTS;
+            constexpr TI PER_AGENT_ACTION_DIM = ACTION_DIM / N_AGENTS;
+            constexpr TI KERNEL_BLOCKSIZE = BATCH_SIZE < 1024 ? (BATCH_SIZE <= 32 ? 32 : (BATCH_SIZE <= 64 ? 64 : (BATCH_SIZE <= 128 ? 128 : (BATCH_SIZE <= 256 ? 256 : (BATCH_SIZE <= 512 ? 512 : 1024))))) : 1024;
+            static_assert(2 * KERNEL_BLOCKSIZE * sizeof(T) <= 48 * 1024, "ppo_per_sample_kernel shared memory usage exceeds 48KB limit");
+            // Fused advantage normalization via shared memory reduction
+            extern __shared__ char shared_mem_raw[];
+            T* shared_sum = reinterpret_cast<T*>(shared_mem_raw);
+            T* shared_sq = shared_sum + blockDim.x;
+            T advantage_mean = 0;
+            T advantage_std = 0;
+            if(PPO_SPEC::PARAMETERS::NORMALIZE_ADVANTAGE){
+                T local_sum = 0, local_sq = 0;
+                for(TI i = threadIdx.x; i < BATCH_SIZE; i += blockDim.x){
+                    T val = get(batch_advantages, i, 0);
+                    local_sum += val;
+                    local_sq += val * val;
+                }
+                shared_sum[threadIdx.x] = local_sum;
+                shared_sq[threadIdx.x] = local_sq;
+                __syncthreads();
+                for(TI s = blockDim.x / 2; s > 0; s >>= 1){
+                    if(threadIdx.x < s){
+                        shared_sum[threadIdx.x] += shared_sum[threadIdx.x + s];
+                        shared_sq[threadIdx.x] += shared_sq[threadIdx.x + s];
+                    }
+                    __syncthreads();
+                }
+                advantage_mean = shared_sum[0] / BATCH_SIZE;
+                T sq_mean = shared_sq[0] / BATCH_SIZE;
+                T variance = sq_mean - advantage_mean * advantage_mean;
+                advantage_std = variance > 0 ? math::sqrt(device.math, variance) : 0;
+            }
+            for(TI batch_step_i = threadIdx.x; batch_step_i < BATCH_SIZE; batch_step_i += blockDim.x){
+                T action_log_prob = 0;
+                for(TI action_i = 0; action_i < ACTION_DIM; action_i++){
+                    T current_action = get(ppo_buffers.current_batch_actions, batch_step_i, action_i);
+                    T rollout_action = get(batch_actions, batch_step_i, action_i);
+                    T current_action_log_std = get(log_std_params, 0, action_i % PER_AGENT_ACTION_DIM);
+                    action_log_prob += random::normal_distribution::log_prob(device.random, current_action, current_action_log_std, rollout_action);
+                    set(ppo_buffers.d_action_log_prob_d_action, batch_step_i, action_i, random::normal_distribution::d_log_prob_d_mean(device.random, current_action, current_action_log_std, rollout_action));
+                    if(PPO_SPEC::PARAMETERS::LEARN_ACTION_STD){
+                        T d_action_log_prob_d_current_action_log_std = random::normal_distribution::d_log_prob_d_log_std(device.random, current_action, current_action_log_std, rollout_action);
+                        set(ppo_buffers.d_action_log_prob_d_action_log_std, batch_step_i, action_i, d_action_log_prob_d_current_action_log_std);
+                    }
+                }
+                T rollout_action_log_prob = get(batch_action_log_probs, batch_step_i, 0);
+                T advantage = get(batch_advantages, batch_step_i, 0);
+                if(PPO_SPEC::PARAMETERS::NORMALIZE_ADVANTAGE){
+                    advantage = (advantage - advantage_mean) / (advantage_std + PPO_SPEC::PARAMETERS::ADVANTAGE_EPSILON);
+                }
+                T log_ratio = action_log_prob - rollout_action_log_prob;
+                T ratio = math::exp(device.math, log_ratio);
+                T clipped_ratio = math::clamp(device.math, ratio, (T)1 - PPO_SPEC::PARAMETERS::EPSILON_CLIP, (T)1 + PPO_SPEC::PARAMETERS::EPSILON_CLIP);
+                bool clipped = ratio != clipped_ratio;
+                T normal_advantage = ratio * advantage;
+                T clipped_advantage = clipped_ratio * advantage;
+                bool ratio_min_switch = normal_advantage - clipped_advantage <= (T)0;
+                T d_loss_d_pessimistic_surrogate = -(T)1/BATCH_SIZE;
+                T d_pessimistic_surrogate_d_normal_advantage = ratio_min_switch ? 1 : 0;
+                T d_pessimistic_surrogate_d_clipped_advantage = ratio_min_switch ? 0 : 1;
+                T d_normal_advantage_d_ratio = advantage;
+                T d_clipped_advantage_d_clipped_ratio = advantage;
+                T d_clipped_ratio_d_ratio = clipped ? 0 : 1;
+                T d_pessimistic_surrogate_d_ratio = d_pessimistic_surrogate_d_normal_advantage * d_normal_advantage_d_ratio + d_pessimistic_surrogate_d_clipped_advantage * d_clipped_advantage_d_clipped_ratio * d_clipped_ratio_d_ratio;
+                T d_loss_d_ratio = d_loss_d_pessimistic_surrogate * d_pessimistic_surrogate_d_ratio;
+                T d_ratio_d_action_log_prob = ratio;
+                T d_loss_d_action_log_prob = d_loss_d_ratio * d_ratio_d_action_log_prob;
+                for(TI action_i = 0; action_i < ACTION_DIM; action_i++){
+                    multiply(ppo_buffers.d_action_log_prob_d_action, batch_step_i, action_i, d_loss_d_action_log_prob);
+                    if(PPO_SPEC::PARAMETERS::LEARN_ACTION_STD){
+                        multiply(ppo_buffers.d_action_log_prob_d_action_log_std, batch_step_i, action_i, d_loss_d_action_log_prob);
+                    }
+                }
+            }
+        }
+        // Deterministic reduction: sum columns of per-sample log_std gradients into log_std_gradient
+        template <typename DEVICE, typename LOG_STD_SPEC, typename D_LOG_STD_SPEC>
+        __global__
+        void reduce_log_std_gradient_kernel(DEVICE device, Matrix<D_LOG_STD_SPEC> d_action_log_prob_d_action_log_std, Matrix<LOG_STD_SPEC> log_std_gradient, typename LOG_STD_SPEC::TI batch_size, typename LOG_STD_SPEC::TI action_dim, typename LOG_STD_SPEC::TI per_agent_action_dim, typename LOG_STD_SPEC::T entropy_grad){
+            using T = typename LOG_STD_SPEC::T;
+            using TI = typename DEVICE::index_t;
+            TI action_i = threadIdx.x + blockIdx.x * blockDim.x;
+            if(action_i < per_agent_action_dim){
+                T sum = 0;
+                TI n_agents = action_dim / per_agent_action_dim;
+                for(TI agent_i = 0; agent_i < n_agents; agent_i++){
+                    TI col = agent_i * per_agent_action_dim + action_i;
+                    for(TI batch_i = 0; batch_i < batch_size; batch_i++){
+                        sum += get(d_action_log_prob_d_action_log_std, batch_i, col);
+                    }
+                }
+                sum += entropy_grad;
+                set(log_std_gradient, 0, action_i, get(log_std_gradient, 0, action_i) + sum);
+            }
+        }
+    }
+    // CUDA overload: GAE
+    template <typename DEV_SPEC, typename DATASET_SPEC, typename PPO_PARAMETERS>
+    void estimate_generalized_advantages(devices::CUDA<DEV_SPEC>& device, rl::components::on_policy_runner::Dataset<DATASET_SPEC>& dataset, PPO_PARAMETERS ppo_parameters_tag){
+        using DEVICE = devices::CUDA<DEV_SPEC>;
+        using TI = typename DEVICE::index_t;
+        using OPR_SPEC = typename DATASET_SPEC::SPEC;
+        constexpr TI BLOCKSIZE = 32;
+        constexpr TI N_BLOCKS = RL_TOOLS_DEVICES_CUDA_CEIL(OPR_SPEC::N_ENVIRONMENTS, BLOCKSIZE);
+        dim3 grid(N_BLOCKS);
+        dim3 block(BLOCKSIZE);
+        devices::cuda::TAG<DEVICE, true> tag_device{};
+        rl::algorithms::ppo::cuda::estimate_generalized_advantages_kernel<<<grid, block, 0, device.stream>>>(tag_device, dataset, ppo_parameters_tag);
+        check_status(device);
+    }
+    // CUDA overload: per-sample PPO loss gradient (with fused advantage normalization)
+    template <typename DEV_SPEC, typename PPO_SPEC, typename BUFFERS_SPEC, typename BATCH_ACTIONS_SPEC, typename BATCH_ACTIONS_MEAN_SPEC, typename BATCH_ACTION_LOG_PROBS_SPEC, typename BATCH_ADVANTAGES_SPEC, typename RNG>
+    void ppo_compute_actor_loss_gradient(devices::CUDA<DEV_SPEC>& device, rl::algorithms::PPO<PPO_SPEC>& ppo, rl::algorithms::ppo::Buffers<BUFFERS_SPEC>& ppo_buffers, Matrix<BATCH_ACTIONS_SPEC>& batch_actions, Matrix<BATCH_ACTIONS_MEAN_SPEC>& batch_actions_mean, Matrix<BATCH_ACTION_LOG_PROBS_SPEC>& batch_action_log_probs, Matrix<BATCH_ADVANTAGES_SPEC>& batch_advantages, typename PPO_SPEC::TYPE_POLICY::DEFAULT& policy_kl_divergence, typename PPO_SPEC::TYPE_POLICY::DEFAULT& batch_policy_kl_divergence, RNG& rng){
+        using DEVICE = devices::CUDA<DEV_SPEC>;
+        using T = typename PPO_SPEC::TYPE_POLICY::DEFAULT;
+        using TI = typename DEVICE::index_t;
+        constexpr TI BATCH_SIZE = PPO_SPEC::PARAMETERS::BATCH_SIZE;
+        constexpr TI ACTION_DIM = PPO_SPEC::ENVIRONMENT::ACTION_DIM;
+        constexpr TI N_AGENTS = PPO_SPEC::ENVIRONMENT::N_AGENTS;
+        constexpr TI PER_AGENT_ACTION_DIM = ACTION_DIM / N_AGENTS;
+        devices::cuda::TAG<DEVICE, true> tag_device{};
+        auto& last_layer = get_last_layer(ppo.actor);
+        auto log_std = matrix_view(device, last_layer.log_std.parameters);
+        auto log_std_grad = matrix_view(device, last_layer.log_std.gradient);
+        // Single block: shared memory reduction computes mean/std, then all threads process their samples
+        // Each thread handles ceil(BATCH_SIZE/BLOCKSIZE) samples for both reduction and per-sample computation
+        constexpr TI KERNEL_BLOCKSIZE = BATCH_SIZE < 1024 ? (BATCH_SIZE <= 32 ? 32 : (BATCH_SIZE <= 64 ? 64 : (BATCH_SIZE <= 128 ? 128 : (BATCH_SIZE <= 256 ? 256 : (BATCH_SIZE <= 512 ? 512 : 1024))))) : 1024;
+        constexpr TI SHARED_MEM_SIZE = 2 * KERNEL_BLOCKSIZE * sizeof(T);
+        static_assert(SHARED_MEM_SIZE <= 48 * 1024, "ppo_per_sample_kernel shared memory exceeds 48KB limit");
+        rl::algorithms::ppo::cuda::ppo_per_sample_kernel<decltype(tag_device), PPO_SPEC, BUFFERS_SPEC, BUFFERS_SPEC><<<1, KERNEL_BLOCKSIZE, SHARED_MEM_SIZE, device.stream>>>(tag_device, ppo_buffers, batch_advantages, batch_actions, batch_action_log_probs, log_std);
+        check_status(device);
+        if(PPO_SPEC::PARAMETERS::LEARN_ACTION_STD){
+            T entropy_grad = -(T)PPO_SPEC::PARAMETERS::ACTION_ENTROPY_COEFFICIENT * (T)N_AGENTS;
+            constexpr TI REDUCE_BLOCKSIZE = 32;
+            constexpr TI REDUCE_N_BLOCKS = RL_TOOLS_DEVICES_CUDA_CEIL(PER_AGENT_ACTION_DIM, REDUCE_BLOCKSIZE);
+            rl::algorithms::ppo::cuda::reduce_log_std_gradient_kernel<<<REDUCE_N_BLOCKS, REDUCE_BLOCKSIZE, 0, device.stream>>>(tag_device, ppo_buffers.d_action_log_prob_d_action_log_std, log_std_grad, (TI)BATCH_SIZE, (TI)ACTION_DIM, (TI)PER_AGENT_ACTION_DIM, entropy_grad);
+            check_status(device);
+        }
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

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

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