hyperdrone 0.1.0__py3-none-any.whl

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/environments/l2f/operations_generic/30_sample_initial_state.h

@@ -0,0 +1,276 @@
+#include "../../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ENVIRONMENTS_L2F_OPERATIONS_GENERIC_SAMPLE_INITIAL_STATE_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ENVIRONMENTS_L2F_OPERATIONS_GENERIC_SAMPLE_INITIAL_STATE_H
+
+#include "../multirotor.h"
+
+#include <rl_tools/utils/generic/vector_operations.h>
+#include "../quaternion_helper.h"
+
+#include <rl_tools/utils/generic/typing.h>
+
+#include <rl_tools/rl/environments/operations_generic.h>
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template<typename DEVICE, typename SPEC, typename PARAMETERS, typename STATE, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT static void sample_initial_state(DEVICE& device, rl::environments::Multirotor<SPEC>& env, PARAMETERS& parameters, STATE& state, RNG& rng);
+    namespace rl::environments::l2f{
+        template <typename DEVICE, typename T, typename RNG>
+        RL_TOOLS_FUNCTION_PLACEMENT void sample_orientation(DEVICE& device, T limit, T output[4], RNG& rng){
+            // Uniform sampling on the desired angle range
+            T u = random::uniform_real_distribution(device.random, (T)0, (T)1, rng);
+            T v = random::uniform_real_distribution(device.random, (T)0, (T)1, rng);
+            T phi = 2.0 * math::PI<T> * u;
+            T cos_theta = 1.0 - 2.0 * v;
+            T sin_theta = math::sqrt(device.math, 1.0 - cos_theta*cos_theta);
+            T x = sin_theta * math::cos(device.math, phi);
+            T y = sin_theta * math::sin(device.math, phi);
+            T z = cos_theta;
+            T angle = random::uniform_real_distribution(device.random, (T)0, limit, rng);
+
+            // Quaternion = [cos(angle/2), sin(angle/2)*axis]
+            T half = 0.5 * angle;
+            T s = math::sin(device.math, half);
+            output[0] = math::cos(device.math, half);
+            output[1] = x * s;
+            output[2] = y * s;
+            output[3] = z * s;
+        }
+        template<typename DEVICE, typename SPEC, typename PARAMETERS, typename STATE_SPEC, typename RNG>
+        RL_TOOLS_FUNCTION_PLACEMENT static void _sample_initial_state(DEVICE& device, Multirotor<SPEC>& env, PARAMETERS& parameters, StateBase<STATE_SPEC>& state, RNG& rng, bool inherited_guidance = false){
+            typename DEVICE::SPEC::MATH math_dev;
+            typename DEVICE::SPEC::RANDOM random_dev;
+            using STATE = StateBase<STATE_SPEC>;
+            using T = typename STATE_SPEC::T;
+            using TI = typename DEVICE::index_t;
+            bool guidance;
+            guidance = random::uniform_real_distribution(random_dev, (T)0, (T)1, rng) < parameters.mdp.init.guidance;
+            if(!guidance){
+                for(TI i = 0; i < 3; i++){
+                    state.position[i] = random::uniform_real_distribution(random_dev, -parameters.mdp.init.max_position, parameters.mdp.init.max_position, rng);
+                }
+            }
+            else{
+                for(TI i = 0; i < 3; i++){
+                    state.position[i] = 0;
+                }
+            }
+            if(parameters.mdp.init.max_angle > 0 && !guidance){
+                sample_orientation(device, parameters.mdp.init.max_angle, state.orientation, rng);
+            }
+            else{
+                state.orientation[0] = 1;
+                state.orientation[1] = 0;
+                state.orientation[2] = 0;
+                state.orientation[3] = 0;
+            }
+            if(!guidance) {
+                for(TI i = 0; i < 3; i++){
+                    state.linear_velocity[i] = random::uniform_real_distribution(random_dev, -parameters.mdp.init.max_linear_velocity, parameters.mdp.init.max_linear_velocity, rng);
+                }
+                for(TI i = 0; i < 3; i++){
+                    state.angular_velocity[i] = random::uniform_real_distribution(random_dev, -parameters.mdp.init.max_angular_velocity, parameters.mdp.init.max_angular_velocity, rng);
+                }
+            }
+            else{
+                for(TI i = 0; i < 3; i++){
+                    state.linear_velocity[i] = 0;
+                }
+                for(TI i = 0; i < 3; i++){
+                    state.angular_velocity[i] = 0;
+                }
+            }
+        }
+        template<typename DEVICE, typename SPEC, typename PARAMETERS, typename STATE_SPEC, typename RNG>
+        RL_TOOLS_FUNCTION_PLACEMENT static void _sample_initial_state(DEVICE& device, Multirotor<SPEC>& env, PARAMETERS& parameters, StateLastAction<STATE_SPEC>& state, RNG& rng){
+            using TI = typename DEVICE::index_t;
+            using STATE = StateLastAction<STATE_SPEC>;
+            sample_initial_state(device, env, parameters, static_cast<typename STATE_SPEC::NEXT_COMPONENT&>(state), rng);
+            for (TI action_i=0; action_i < STATE::ACTION_DIM; action_i++){
+                state.last_action[action_i] = 0;
+            }
+        }
+        template<typename DEVICE, typename SPEC, typename PARAMETERS, typename STATE_SPEC, typename RNG>
+        RL_TOOLS_FUNCTION_PLACEMENT static void _sample_initial_state(DEVICE& device, Multirotor<SPEC>& env, PARAMETERS& parameters, StateLinearAcceleration<STATE_SPEC>& state, RNG& rng){
+            using TI = typename DEVICE::index_t;
+            sample_initial_state(device, env, parameters, static_cast<typename STATE_SPEC::NEXT_COMPONENT&>(state), rng);
+            for(TI i = 0; i < 3; i++){
+                state.linear_acceleration[i] = 0;
+            }
+        }
+        template<typename DEVICE, typename SPEC, typename PARAMETERS, typename STATE_SPEC, typename RNG>
+        RL_TOOLS_FUNCTION_PLACEMENT static void _sample_initial_state(DEVICE& device, Multirotor<SPEC>& env, PARAMETERS& parameters, StateGyroBias<STATE_SPEC>& state, RNG& rng){
+            typename DEVICE::SPEC::RANDOM random_dev;
+            using T = typename SPEC::T;
+            using TI = typename DEVICE::index_t;
+            sample_initial_state(device, env, parameters, static_cast<typename STATE_SPEC::NEXT_COMPONENT&>(state), rng);
+            T limit = parameters.imu.gyro_bias.init_max;
+            for(TI i = 0; i < 3; i++){
+                state.gyro_bias[i] = random::uniform_real_distribution(random_dev, -limit, limit, rng);
+            }
+        }
+        template<typename DEVICE, typename SPEC, typename PARAMETERS, typename STATE_SPEC, typename RNG>
+        RL_TOOLS_FUNCTION_PLACEMENT static void _sample_initial_state(DEVICE& device, Multirotor<SPEC>& env, PARAMETERS& parameters, StateMahony<STATE_SPEC>& state, RNG& rng){
+            using TI = typename DEVICE::index_t;
+            using T = typename STATE_SPEC::T;
+            sample_initial_state(device, env, parameters, static_cast<typename STATE_SPEC::NEXT_COMPONENT&>(state), rng);
+            state.q_estimate[0] = 1;
+            state.q_estimate[1] = 0;
+            state.q_estimate[2] = 0;
+            state.q_estimate[3] = 0;
+            T conjugate_orientation[4];
+            conjugate_orientation[0] =  state.orientation[0];
+            conjugate_orientation[1] = -state.orientation[1];
+            conjugate_orientation[2] = -state.orientation[2];
+            conjugate_orientation[3] = -state.orientation[3];
+            T accel_global[3];
+            accel_global[0] = -parameters.dynamics.gravity[0];
+            accel_global[1] = -parameters.dynamics.gravity[1];
+            accel_global[2] = -parameters.dynamics.gravity[2];
+            T accel_body[3];
+            rotate_vector_by_quaternion<DEVICE, T>(conjugate_orientation, accel_global, accel_body);
+            mahony_quaternion_from_accel(device, accel_body, state.q_estimate);
+            for(TI i = 0; i < 3; i++){
+                state.bias_estimate[i] = 0;
+            }
+        }
+        template<typename DEVICE, typename SPEC, typename PARAMETERS, typename STATE_SPEC, typename RNG>
+        RL_TOOLS_FUNCTION_PLACEMENT static void _sample_initial_state(DEVICE& device, Multirotor<SPEC>& env, PARAMETERS& parameters, StateCTBRController<STATE_SPEC>& state, RNG& rng){
+            using TI = typename DEVICE::index_t;
+            sample_initial_state(device, env, parameters, static_cast<typename STATE_SPEC::NEXT_COMPONENT&>(state), rng);
+            for(TI i = 0; i < 3; i++){
+                state.previous_angular_velocity[i] = state.angular_velocity[i];
+            }
+        }
+        template<typename DEVICE, typename SPEC, typename PARAMETERS, typename STATE_SPEC, typename RNG>
+        RL_TOOLS_FUNCTION_PLACEMENT static void _sample_initial_state(DEVICE& device, Multirotor<SPEC>& env, PARAMETERS& parameters, StateLinearAccelerationHistory<STATE_SPEC>& state, RNG& rng){
+            using TI = typename DEVICE::index_t;
+            using STATE = StateLinearAccelerationHistory<STATE_SPEC>;
+            sample_initial_state(device, env, parameters, static_cast<typename STATE_SPEC::NEXT_COMPONENT&>(state), rng);
+            state.acceleration_history_step = 0;
+            for(TI step_i = 0; step_i < STATE::HISTORY_MEM_LENGTH; step_i++){
+                for(TI dim_i = 0; dim_i < STATE::ACCELERATION_DIM; dim_i++){
+                    state.linear_acceleration_body_history[step_i][dim_i] = 0;
+                }
+            }
+        }
+        template<typename DEVICE, typename SPEC, typename PARAMETERS, typename STATE_SPEC, typename RNG>
+        RL_TOOLS_FUNCTION_PLACEMENT static void _sample_initial_state(DEVICE& device, Multirotor<SPEC>& env, PARAMETERS& parameters, StateAngularVelocityDelay<STATE_SPEC>& state, RNG& rng){
+            using TI = typename DEVICE::index_t;
+            using STATE = StateAngularVelocityDelay<STATE_SPEC>;
+            sample_initial_state(device, env, parameters, static_cast<typename STATE_SPEC::NEXT_COMPONENT&>(state), rng);
+            for (TI step_i=0; step_i < STATE::HISTORY_MEM_LENGTH; step_i++){
+                for (TI dim_i=0; dim_i < 3; dim_i++){
+                    state.angular_velocity_history[step_i][dim_i] = state.angular_velocity[dim_i];
+                }
+            }
+        }
+        template<typename DEVICE, typename SPEC, typename PARAMETERS, typename STATE_SPEC, typename RNG>
+        RL_TOOLS_FUNCTION_PLACEMENT static void _sample_initial_state(DEVICE& device, Multirotor<SPEC>& env, PARAMETERS& parameters, StateLinearVelocityDelay<STATE_SPEC>& state, RNG& rng){
+            using TI = typename DEVICE::index_t;
+            using STATE = StateLinearVelocityDelay<STATE_SPEC>;
+            sample_initial_state(device, env, parameters, static_cast<typename STATE_SPEC::NEXT_COMPONENT&>(state), rng);
+            for (TI step_i=0; step_i < STATE::HISTORY_MEM_LENGTH; step_i++){
+                for (TI dim_i=0; dim_i < 3; dim_i++){
+                    state.linear_velocity_history[step_i][dim_i] = state.linear_velocity[dim_i];
+                }
+            }
+        }
+        template<typename DEVICE, typename SPEC, typename PARAMETERS, typename STATE_SPEC, typename RNG>
+        RL_TOOLS_FUNCTION_PLACEMENT static void _sample_initial_state(DEVICE& device, Multirotor<SPEC>& env, PARAMETERS& parameters, StatePoseErrorIntegral<STATE_SPEC>& state, RNG& rng){
+            using TI = typename DEVICE::index_t;
+            sample_initial_state(device, env, parameters, static_cast<typename STATE_SPEC::NEXT_COMPONENT&>(state), rng);
+            for (TI dim_i=0; dim_i<3; dim_i++){
+                state.position_integral[dim_i] = 0;
+            }
+            state.orientation_integral = 0;
+        }
+        template<typename DEVICE, typename SPEC, typename PARAMETERS, typename STATE_SPEC, typename RNG>
+        RL_TOOLS_FUNCTION_PLACEMENT static void _sample_initial_state(DEVICE& device, Multirotor<SPEC>& env, PARAMETERS& parameters, StateRandomForce<STATE_SPEC>& state, RNG& rng){
+            typename DEVICE::SPEC::RANDOM random_dev;
+            using T = typename SPEC::T;
+    //        bool guidance = random::uniform_real_distribution(random_dev, (T)0, (T)1, rng) < parameters.mdp.init.guidance;
+            sample_initial_state(device, env, parameters, static_cast<typename STATE_SPEC::NEXT_COMPONENT&>(state), rng);
+    //        if(!guidance){
+            {
+                auto distribution = parameters.disturbances.random_force;
+                state.force[0] = random::normal_distribution::sample(random_dev, (T)distribution.mean, (T)distribution.std, rng);
+                state.force[1] = random::normal_distribution::sample(random_dev, (T)distribution.mean, (T)distribution.std, rng);
+                state.force[2] = random::normal_distribution::sample(random_dev, (T)distribution.mean, (T)distribution.std, rng);
+            }
+            {
+                auto distribution = parameters.disturbances.random_torque;
+                state.torque[0] = random::normal_distribution::sample(random_dev, (T)distribution.mean, (T)distribution.std, rng);
+                state.torque[1] = random::normal_distribution::sample(random_dev, (T)distribution.mean, (T)distribution.std, rng);
+                state.torque[2] = random::normal_distribution::sample(random_dev, (T)distribution.mean, (T)distribution.std/100, rng);
+            }
+    //        }
+    //        else{
+    //            state.force[0] = 0;
+    //            state.force[1] = 0;
+    //            state.force[2] = 0;
+    //            state.torque[0] = 0;
+    //            state.torque[1] = 0;
+    //            state.torque[2] = 0;
+    //        }
+
+        }
+        template<typename DEVICE, typename SPEC, typename PARAMETERS, typename STATE_SPEC, typename RNG>
+        RL_TOOLS_FUNCTION_PLACEMENT static void _sample_initial_state(DEVICE& device, Multirotor<SPEC>& env, PARAMETERS& parameters, StateRandomOrientationOffset<STATE_SPEC>& state, RNG& rng){
+            typename DEVICE::SPEC::RANDOM random_dev;
+            using T = typename SPEC::T;
+            sample_initial_state(device, env, parameters, static_cast<typename STATE_SPEC::NEXT_COMPONENT&>(state), rng);
+            sample_orientation(device, parameters.domain_randomization, state.orientation, rng);
+        }
+        template<typename DEVICE, typename SPEC, typename PARAMETERS, typename STATE_SPEC, typename RNG>
+        RL_TOOLS_FUNCTION_PLACEMENT static void _sample_initial_state(DEVICE& device, Multirotor<SPEC>& env, PARAMETERS& parameters, StateRotors<STATE_SPEC>& state, RNG& rng){
+            sample_initial_state(device, env, parameters, static_cast<typename STATE_SPEC::NEXT_COMPONENT&>(state), rng);
+            using TI = typename DEVICE::index_t;
+            using T = typename SPEC::T;
+            T min_rpm, max_rpm;
+            if(parameters.mdp.init.relative_rpm){
+                min_rpm = (parameters.mdp.init.min_rpm + 1)/2 * (parameters.dynamics.action_limit.max - parameters.dynamics.action_limit.min) + parameters.dynamics.action_limit.min;
+                max_rpm = (parameters.mdp.init.max_rpm + 1)/2 * (parameters.dynamics.action_limit.max - parameters.dynamics.action_limit.min) + parameters.dynamics.action_limit.min;
+            }
+            else{
+                min_rpm = parameters.mdp.init.min_rpm < 0 ? parameters.dynamics.action_limit.min : parameters.mdp.init.min_rpm;
+                max_rpm = parameters.mdp.init.max_rpm < 0 ? parameters.dynamics.action_limit.max : parameters.mdp.init.max_rpm;
+                if(max_rpm > parameters.dynamics.action_limit.max){
+                    max_rpm = parameters.dynamics.action_limit.max;
+                }
+                if(min_rpm > max_rpm){
+                    min_rpm = max_rpm;
+                }
+            }
+            for(TI i = 0; i < 4; i++){
+                state.rpm[i] = random::uniform_real_distribution(typename DEVICE::SPEC::RANDOM(), min_rpm, max_rpm, rng);
+            }
+        }
+        template<typename DEVICE, typename SPEC, typename PARAMETERS, typename STATE_SPEC, typename RNG>
+        RL_TOOLS_FUNCTION_PLACEMENT static void _sample_initial_state(DEVICE& device, Multirotor<SPEC>& env, PARAMETERS& parameters, StateRotorsHistory<STATE_SPEC>& state, RNG& rng){
+            using MULTIROTOR = Multirotor<SPEC>;
+            using TI = typename DEVICE::index_t;
+            using STATE = StateRotorsHistory<STATE_SPEC>;
+            sample_initial_state(device, env, parameters, static_cast<typename STATE::NEXT_COMPONENT&>(state), rng);
+            state.rotor_history_step = 0;
+            for(TI step_i = 0; step_i < STATE_SPEC::HISTORY_LENGTH; step_i++){
+                for(TI action_i = 0; action_i < MULTIROTOR::ACTION_DIM; action_i++){
+                    state.action_history[step_i][action_i] = (state.rpm[action_i] - parameters.dynamics.action_limit.min) / (parameters.dynamics.action_limit.max - parameters.dynamics.action_limit.min) * 2 - 1;
+                }
+            }
+        }
+        template<typename DEVICE, typename SPEC, typename PARAMETERS, typename STATE_SPEC, typename RNG>
+        RL_TOOLS_FUNCTION_PLACEMENT static void _sample_initial_state(DEVICE& device, Multirotor<SPEC>& env, PARAMETERS& parameters, StateTrajectory<STATE_SPEC>& state, RNG& rng){
+            using MULTIROTOR = Multirotor<SPEC>;
+            using TI = typename DEVICE::index_t;
+            using T = typename SPEC::T;
+            using STATE = StateTrajectory<STATE_SPEC>;
+            sample_initial_state(device, env, parameters, static_cast<typename STATE::NEXT_COMPONENT&>(state), rng);
+            state.trajectory_step = 0;
+        }
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/environments/l2f/operations_generic/35_get_desired_state.h

@@ -0,0 +1,63 @@
+#include "../../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ENVIRONMENTS_L2F_OPERATIONS_GENERIC_GET_DESIRED_STATE_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ENVIRONMENTS_L2F_OPERATIONS_GENERIC_GET_DESIRED_STATE_H
+
+#include "../multirotor.h"
+
+#include <rl_tools/utils/generic/vector_operations.h>
+#include "../quaternion_helper.h"
+
+#include <rl_tools/utils/generic/typing.h>
+
+#include <rl_tools/rl/environments/operations_generic.h>
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template<typename DEVICE, typename SPEC, typename PARAMETERS, typename STATE_SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT static void get_desired_state(DEVICE& device, const rl::environments::Multirotor<SPEC>& env, PARAMETERS& parameters, const rl::environments::l2f::StateBase<STATE_SPEC>& state, rl::environments::l2f::StateBase<STATE_SPEC>& desired_state, RNG& rng){
+        using TI = typename DEVICE::index_t;
+        using T = typename SPEC::T;
+        desired_state.position[0] = 0;
+        desired_state.position[1] = 0;
+        desired_state.position[2] = 0;
+        desired_state.linear_velocity[0] = 0;
+        desired_state.linear_velocity[1] = 0;
+        desired_state.linear_velocity[2] = 0;
+    }
+    template<typename T_TI>
+    struct TrajectoryIndexResult{
+        T_TI index;
+        bool forward;
+    };
+    template<typename DEVICE, typename SPEC, typename PARAMETERS, typename STATE_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT TrajectoryIndexResult<typename DEVICE::index_t> get_trajectory_index(DEVICE& device, const rl::environments::Multirotor<SPEC>& env, PARAMETERS& parameters, const rl::environments::l2f::StateTrajectory<STATE_SPEC>& state, typename DEVICE::index_t offset=0){
+        using TI = typename DEVICE::index_t;
+        using ENVIRONMENT = rl::environments::Multirotor<SPEC>;
+        static constexpr TI EPISODE_STEP_LIMIT = ENVIRONMENT::EPISODE_STEP_LIMIT;
+        static_assert(decltype(parameters.trajectory)::LENGTH == EPISODE_STEP_LIMIT);
+        TI full_step = state.trajectory_step + offset;
+        TI interval = full_step / EPISODE_STEP_LIMIT;
+        bool forward = interval % 2 == 0;
+        TI progress = full_step % EPISODE_STEP_LIMIT;
+        TI index = forward ? progress : (EPISODE_STEP_LIMIT - progress - 1);
+        return {index, forward};
+    }
+    template<typename DEVICE, typename SPEC, typename PARAMETERS, typename STATE_SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT static void get_desired_state(DEVICE& device, const rl::environments::Multirotor<SPEC>& env, PARAMETERS& parameters, const rl::environments::l2f::StateTrajectory<STATE_SPEC>& state, rl::environments::l2f::StateTrajectory<STATE_SPEC>& desired_state, RNG& rng){
+        using TI = typename DEVICE::index_t;
+        using T = typename SPEC::T;
+        T traj_dt = (T)decltype(parameters.trajectory)::DT;
+        utils::assert_exit(device, parameters.integration.dt == traj_dt, "Mismatch between environment and trajectory parameters integration dt");
+        auto traj_idx = get_trajectory_index(device, env, parameters, state);
+        T direction = traj_idx.forward ? (T)1 : (T)-1;
+        for (TI axis_i=0; axis_i < 3; axis_i++){
+            desired_state.position[axis_i]        = parameters.trajectory.steps[traj_idx.index].position[axis_i];
+            desired_state.linear_velocity[axis_i] = parameters.trajectory.steps[traj_idx.index].linear_velocity[axis_i] * direction;
+        }
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif
+
+