hyperdrone 0.1.0__py3-none-any.whl

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

@@ -0,0 +1,31 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_COMPONENTS_RUNNING_NORMALIZER_PERSIST_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_COMPONENTS_RUNNING_NORMALIZER_PERSIST_H
+#include "running_normalizer.h"
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template<typename DEVICE, typename SPEC, typename GROUP>
+    void save(DEVICE& device, rl::components::RunningNormalizer<SPEC>& normalizer, GROUP& group) {
+        save(device, normalizer.mean, group, "mean");
+        save(device, normalizer.std, group, "std");
+        Tensor<tensor::Specification<typename SPEC::TI, typename SPEC::TI, tensor::Shape<typename SPEC::TI, 1>>> age_tensor;
+        malloc(device, age_tensor);
+        set(device, age_tensor, normalizer.age, 0);
+        save(device, age_tensor, group, "age");
+        free(device, age_tensor);
+    }
+    template<typename DEVICE, typename SPEC, typename GROUP>
+    bool load(DEVICE& device, rl::components::RunningNormalizer<SPEC>& normalizer, GROUP& group) {
+        bool success = load(device, normalizer.mean, group, "mean");
+        success &= load(device, normalizer.std, group, "std");
+        Tensor<tensor::Specification<typename SPEC::TI, typename SPEC::TI, tensor::Shape<typename SPEC::TI, 1>>> age_tensor;
+        malloc(device, age_tensor);
+        success &= load(device, age_tensor, group, "age");
+        normalizer.age = get(device, age_tensor, 0);
+        free(device, age_tensor);
+        return success;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

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

@@ -0,0 +1,33 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_COMPONENTS_RUNNING_NORMALIZER_RUNNING_NORMALIZER_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_COMPONENTS_RUNNING_NORMALIZER_RUNNING_NORMALIZER_H
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::rl::components{
+    namespace running_normalizer{
+        template <typename T_TYPE_POLICY, typename T_TI, T_TI T_DIM, bool T_DYNAMIC_ALLOCATION=true>
+        struct Specification{
+            using TYPE_POLICY = T_TYPE_POLICY;
+            using TI = T_TI;
+            static constexpr TI DIM = T_DIM;
+            static constexpr bool DYNAMIC_ALLOCATION = T_DYNAMIC_ALLOCATION;
+        };
+    }
+    template <typename T_SPEC>
+    struct RunningNormalizer{
+        using SPEC = T_SPEC;
+        using TYPE_POLICY = typename SPEC::TYPE_POLICY;
+        using TI = typename SPEC::TI;
+        static constexpr TI DIM = SPEC::DIM;
+        static constexpr bool DYNAMIC_ALLOCATION = SPEC::DYNAMIC_ALLOCATION;
+
+        using T = typename TYPE_POLICY::DEFAULT;
+        Matrix<matrix::Specification<T, TI, 1, DIM, DYNAMIC_ALLOCATION>> mean;
+        Matrix<matrix::Specification<T, TI, 1, DIM, DYNAMIC_ALLOCATION>> std;
+        TI age = 0;
+    };
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif

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

@@ -0,0 +1,58 @@
+#include "../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ENVIRONMENT_WRAPPERS_OPERATIONS_GENERIC_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ENVIRONMENT_WRAPPERS_OPERATIONS_GENERIC_H
+#include "wrappers.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template<typename DEVICE, typename ENVIRONMENT>
+    RL_TOOLS_FUNCTION_PLACEMENT static void malloc(DEVICE& device, rl::environment_wrappers::Wrapper<ENVIRONMENT>& env){
+        malloc(device, env.env);
+    }
+    template<typename DEVICE, typename ENVIRONMENT>
+    RL_TOOLS_FUNCTION_PLACEMENT static void free(DEVICE& device, rl::environment_wrappers::Wrapper<ENVIRONMENT>& env){
+        free(device, env.env);
+    }
+    template<typename DEVICE, typename ENVIRONMENT>
+    RL_TOOLS_FUNCTION_PLACEMENT static void init(DEVICE& device, rl::environment_wrappers::Wrapper<ENVIRONMENT>& env){
+        init(device, env.env);
+    }
+    template<typename DEVICE, typename ENVIRONMENT, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT static void sample_initial_parameters(DEVICE& device, const rl::environment_wrappers::Wrapper<ENVIRONMENT>& env, typename rl::environment_wrappers::Wrapper<ENVIRONMENT>::Parameters& parameters, RNG& rng){
+        sample_initial_parameters(device, env.env, parameters, rng);
+    }
+    template<typename DEVICE, typename ENVIRONMENT>
+    static void initial_parameters(DEVICE& device, const rl::environment_wrappers::Wrapper<ENVIRONMENT>& env, typename rl::environment_wrappers::Wrapper<ENVIRONMENT>::Parameters& parameters){
+        initial_parameters(device, env.env, parameters);
+    }
+    template<typename DEVICE, typename ENVIRONMENT, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT static void sample_initial_state(DEVICE& device, const rl::environment_wrappers::Wrapper<ENVIRONMENT>& env, typename rl::environment_wrappers::Wrapper<ENVIRONMENT>::Parameters& parameters, typename rl::environment_wrappers::Wrapper<ENVIRONMENT>::State& state, RNG& rng){
+        sample_initial_state(device, env.env, parameters, state, rng);
+    }
+    template<typename DEVICE, typename ENVIRONMENT>
+    static void initial_state(DEVICE& device, const rl::environment_wrappers::Wrapper<ENVIRONMENT>& env, typename rl::environment_wrappers::Wrapper<ENVIRONMENT>::Parameters& parameters, typename rl::environment_wrappers::Wrapper<ENVIRONMENT>::State& state){
+        initial_state(device, env.env, parameters, state);
+    }
+    template<typename DEVICE, typename ENVIRONMENT, typename ACTION_SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT typename ENVIRONMENT::T step(DEVICE& device, const rl::environment_wrappers::Wrapper<ENVIRONMENT>& env, typename rl::environment_wrappers::Wrapper<ENVIRONMENT>::Parameters& parameters, const typename rl::environment_wrappers::Wrapper<ENVIRONMENT>::State& state, const Matrix<ACTION_SPEC>& action, typename rl::environment_wrappers::Wrapper<ENVIRONMENT>::State& next_state, RNG& rng) {
+        static_assert(ACTION_SPEC::ROWS == 1);
+        return step(device, env.env, parameters, state, action, next_state, rng);
+    }
+    template<typename DEVICE, typename ENVIRONMENT, typename ACTION_SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT static typename ENVIRONMENT::T reward(DEVICE& device, const rl::environment_wrappers::Wrapper<ENVIRONMENT>& env, typename rl::environment_wrappers::Wrapper<ENVIRONMENT>::Parameters& parameters, const typename rl::environment_wrappers::Wrapper<ENVIRONMENT>::State& state, const Matrix<ACTION_SPEC>& action, const typename rl::environment_wrappers::Wrapper<ENVIRONMENT>::State& next_state, RNG& rng){
+        return reward(device, env.env, parameters, state, action, next_state, rng);
+    }
+
+    template<typename DEVICE, typename ENVIRONMENT, typename OBS_TYPE, typename OBS_SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT static void observe(DEVICE& device, const rl::environment_wrappers::Wrapper<ENVIRONMENT>& env, const typename rl::environment_wrappers::Wrapper<ENVIRONMENT>::Parameters& parameters, const typename rl::environment_wrappers::Wrapper<ENVIRONMENT>::State& state, OBS_TYPE& obs_type, Matrix<OBS_SPEC>& observation, RNG& rng){
+        static_assert(OBS_SPEC::ROWS == 1);
+        observe(device, env.env, parameters, state, obs_type, observation, rng);
+    }
+    template<typename DEVICE, typename ENVIRONMENT, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT static bool terminated(DEVICE& device, const rl::environment_wrappers::Wrapper<ENVIRONMENT>& env, typename rl::environment_wrappers::Wrapper<ENVIRONMENT>::Parameters& parameters, const typename rl::environment_wrappers::Wrapper<ENVIRONMENT>::State state, RNG& rng){
+        return terminated(device, env.env, parameters, state, rng);
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/environment_wrappers/scale_observations/operations_generic.h

@@ -0,0 +1,28 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ENVIRONMENT_WRAPPERS_SCALE_OBSERVATIONS_OPERATIONS_GENERIC_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ENVIRONMENT_WRAPPERS_SCALE_OBSERVATIONS_OPERATIONS_GENERIC_H
+
+#include "../../../math/operations_generic.h"
+#include "wrapper.h"
+#include "../operations_generic.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template<typename DEVICE, typename SPEC, typename ENVIRONMENT, typename OBS_TYPE, typename OBS_SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT static void observe(DEVICE& device, const rl::environment_wrappers::ScaleObservations<SPEC, ENVIRONMENT>& env, const typename rl::environment_wrappers::ScaleObservations<SPEC, ENVIRONMENT>::Parameters& parameters, const typename rl::environment_wrappers::ScaleObservations<SPEC, ENVIRONMENT>::State& state, const OBS_TYPE& obs_type, Matrix<OBS_SPEC>& observation, RNG& rng){
+        using T = typename SPEC::TYPE_POLICY::DEFAULT;
+        using TI = typename DEVICE::index_t;
+        static_assert(OBS_SPEC::ROWS == 1);
+        observe(device, env.env, parameters, state, obs_type, observation, rng);
+        for(TI observation_i = 0; observation_i < OBS_SPEC::COLS; observation_i++){
+            T observation_value = get(observation, 0, observation_i);
+            set(observation, 0, observation_i, observation_value * SPEC::SCALE);
+        }
+    }
+}
+
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif
+

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/environment_wrappers/scale_observations/wrapper.h

@@ -0,0 +1,26 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ENVIRONMENT_WRAPPERS_SCALE_OBSERVATIONS_WRAPPER_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ENVIRONMENT_WRAPPERS_SCALE_OBSERVATIONS_WRAPPER_H
+
+#include "../wrappers.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::rl::environment_wrappers::scale_observations {
+    template <typename T_TYPE_POLICY, typename T_TI>
+    struct Specification{
+        using TYPE_POLICY = T_TYPE_POLICY;
+        using TI = T_TI;
+        static constexpr typename TYPE_POLICY::DEFAULT SCALE = 1;
+    };
+}
+namespace rl_tools::rl::environment_wrappers{
+    template <typename T_SPEC, typename T_ENVIRONMENT>
+    struct ScaleObservations: Wrapper<T_ENVIRONMENT>{
+        using SPEC = T_SPEC;
+    };
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif
+

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/environment_wrappers/wrappers.h

@@ -0,0 +1,29 @@
+#include "../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ENVIRONMENT_WRAPPERS_WRAPPER_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ENVIRONMENT_WRAPPERS_WRAPPER_H
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::rl::environment_wrappers {
+    template <typename T_ENVIRONMENT>
+    struct Wrapper{
+        using ENVIRONMENT = T_ENVIRONMENT;
+        using T = typename ENVIRONMENT::T;
+        using TI = typename ENVIRONMENT::TI;
+        using Observation = typename ENVIRONMENT::Observation;
+        using ObservationPrivileged = typename ENVIRONMENT::ObservationPrivileged;
+        static constexpr TI N_AGENTS = ENVIRONMENT::N_AGENTS;
+        static constexpr TI ACTION_DIM = ENVIRONMENT::ACTION_DIM;
+        static constexpr TI EPISODE_STEP_LIMIT = ENVIRONMENT::EPISODE_STEP_LIMIT;
+        using Parameters = typename ENVIRONMENT::Parameters;
+        using State = typename ENVIRONMENT::State;
+
+        ENVIRONMENT env;
+    };
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#include "scale_observations/wrapper.h"
+
+#endif
+

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/environments/acrobot/acrobot.h

@@ -0,0 +1,88 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ENVIRONMENTS_ACROBOT_ACROBOT_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ENVIRONMENTS_ACROBOT_ACROBOT_H
+
+#include "../environments.h"
+#include "../observation.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::rl::environments::acrobot {
+    template <typename T>
+    struct DefaultParameters {
+        static constexpr T DT = 0.2;
+        static constexpr T LINK_LENGTH_1 = 1.0;  // [m]
+        static constexpr T LINK_LENGTH_2 = 1.0;  // [m]
+        static constexpr T LINK_MASS_1 = 1.0;  //: [kg] mass of link 1
+        static constexpr T LINK_MASS_2 = 1.0;  //: [kg] mass of link 2
+        static constexpr T LINK_COM_POS_1 = 0.5;  //: [m] position of the center of mass of link 1
+        static constexpr T LINK_COM_POS_2 = 0.5;  //: [m] position of the center of mass of link 2
+        static constexpr T LINK_MOI = 1.0;  //: moments of inertia for both links
+        static constexpr T MAX_VEL_1 = 4 * math::PI<T>;
+        static constexpr T MAX_VEL_2 = 9 * math::PI<T>;
+        static constexpr T MIN_TORQUE = -5;
+        static constexpr T MAX_TORQUE = +5;
+        static constexpr T TORQUE_NOISE_MAX = 0.0;
+        static constexpr T SCREEN_DIM = 500;
+        static constexpr T VEL_PENALTY = 0.2;
+    };
+    template <typename T>
+    struct EasyParameters: DefaultParameters<T> {
+        static constexpr T MIN_TORQUE = -10;
+        static constexpr T MAX_TORQUE = +10;
+    };
+    template <typename T>
+    struct Parameters{};
+    template <typename T_T, typename T_TI, typename T_PARAMETERS = DefaultParameters<T_T>>
+    struct Specification{
+        using T = T_T;
+        using TI = T_TI;
+        using PARAMETERS = T_PARAMETERS;
+    };
+
+    template <typename T, typename TI>
+    struct State{
+        static constexpr TI DIM = 4;
+        T theta_1;
+        T theta_2;
+        T theta_1_dot;
+        T theta_2_dot;
+    };
+
+    template <typename TI>
+    struct Observation{
+        static constexpr TI DIM = 6;
+        using SHAPE = tensor::Shape<TI, DIM>;
+    };
+
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::rl::environments{
+
+    template <typename T_SPEC>
+    struct Acrobot: Environment<typename T_SPEC::T, typename T_SPEC::TI> {
+        using SPEC = T_SPEC;
+        using T = typename SPEC::T;
+        using TI = typename SPEC::TI;
+        using State = acrobot::State<T, TI>;
+        using Parameters = typename SPEC::PARAMETERS;
+        using Observation = acrobot::Observation<TI>;
+        using ObservationPrivileged = Observation;
+        static constexpr TI ACTION_DIM = 1;
+        acrobot::Parameters<T> parameters;
+        static constexpr TI EPISODE_STEP_LIMIT = 500;
+    };
+    template <typename T_SPEC>
+    struct AcrobotSwingup: Acrobot<T_SPEC>{};
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+
+
+
+
+
+
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/environments/acrobot/operations_cpu.h

@@ -0,0 +1,152 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ENVIRONMENTS_ACROBOT_OPERATIONS_CPU_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ENVIRONMENTS_ACROBOT_OPERATIONS_CPU_H
+
+#include "acrobot.h"
+#include "operations_generic.h"
+
+#include <string>
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template <typename DEVICE, typename SPEC>
+    std::string json(DEVICE&, const rl::environments::Acrobot<SPEC>& env){
+        std::string json = "{";
+        json += "}";
+        return json;
+    }
+    template <typename DEVICE, typename SPEC>
+    std::string json(DEVICE& device, const rl::environments::Acrobot<SPEC>& env, const typename rl::environments::Acrobot<SPEC>::Parameters& parameters){
+        using PARAMS = typename rl::environments::Acrobot<SPEC>::Parameters;
+        std::string json = "{\"PENDULUM_LENGTH_1\": " + std::to_string(PARAMS::LINK_LENGTH_1) + "," +
+                            "\"PENDULUM_LENGTH_2\": " + std::to_string(PARAMS::LINK_LENGTH_2) + "," +
+                            "\"BOB_RADIUS\": 0.10, \"PIVOT_RADIUS\": 0.05, \"ROD_WIDTH\": 0.04" + "}";
+        return json;
+    }
+    template <typename DEVICE, typename SPEC>
+    std::string json(DEVICE&, const rl::environments::Acrobot<SPEC>& env, const typename rl::environments::Acrobot<SPEC>::Parameters& parameters, const typename rl::environments::Acrobot<SPEC>::State& state){
+        std::string json = "{";
+        json += "\"theta1\":" + std::to_string(state.theta_1) + ",";
+        json += "\"theta2\":" + std::to_string(state.theta_2) + ",";
+        json += "\"theta1_dot\":" + std::to_string(state.theta_1_dot) + ",";
+        json += "\"theta2_dot\":" + std::to_string(state.theta_2_dot);
+        json += "}";
+        return json;
+    }
+    template <typename DEVICE, typename SPEC>
+    std::string json(DEVICE& device, const rl::environments::AcrobotSwingup<SPEC>& env, const typename rl::environments::Acrobot<SPEC>::Parameters& parameters, const typename rl::environments::Acrobot<SPEC>::State& state){
+        return json(device, static_cast<const rl::environments::Acrobot<SPEC>&>(env), parameters, state);
+    }
+
+    template <typename DEVICE, typename SPEC>
+    std::string get_ui(DEVICE& device, rl::environments::Acrobot<SPEC>& env){
+        std::string ui = R"RL_TOOLS_LITERAL(
+export async function init(canvas, options){
+    // Simply saving the context for 2D environments
+    return {
+        ctx: canvas.getContext('2d')
+    }
+}
+export async function render(ui_state, parameters, state, action) {
+    const ctx = ui_state.ctx
+    ctx.clearRect(0, 0, ctx.canvas.width, ctx.canvas.height);
+
+    const centerX = ctx.canvas.width / 2;
+    const centerY = ctx.canvas.height / 2;
+    const canvasWidth = ctx.canvas.width/5;
+
+    const pendulumLength1 = canvasWidth * parameters.PENDULUM_LENGTH_1;
+    const pendulumLength2 = canvasWidth * parameters.PENDULUM_LENGTH_2;
+    const bobRadius = canvasWidth * parameters.BOB_RADIUS;
+    const pivotRadius = canvasWidth * parameters.PIVOT_RADIUS;
+
+    // First pendulum (base to joint)
+    const pendulumX1 = centerX + pendulumLength1 * Math.sin(state.theta1);
+    const pendulumY1 = centerY + pendulumLength1 * Math.cos(state.theta1);
+
+    ctx.beginPath();
+    ctx.moveTo(centerX, centerY);
+    ctx.lineTo(pendulumX1, pendulumY1);
+    ctx.lineWidth = canvasWidth * parameters.ROD_WIDTH;
+    ctx.strokeStyle = 'black';
+    ctx.stroke();
+
+    ctx.beginPath();
+    ctx.arc(centerX, centerY, pivotRadius, 0, 2 * Math.PI);
+    ctx.fillStyle = 'black';
+    ctx.fill();
+    ctx.stroke();
+
+    // Second pendulum (joint to end)
+    const pendulumX2 = pendulumX1 + pendulumLength2 * Math.sin(state.theta1 + state.theta2);
+    const pendulumY2 = pendulumY1 + pendulumLength2 * Math.cos(state.theta1 + state.theta2);
+
+    ctx.beginPath();
+    ctx.moveTo(pendulumX1, pendulumY1);
+    ctx.lineTo(pendulumX2, pendulumY2);
+    ctx.lineWidth = canvasWidth * parameters.ROD_WIDTH;
+    ctx.strokeStyle = 'black';
+    ctx.stroke();
+
+    ctx.lineWidth = 0.04 * canvasWidth
+    ctx.beginPath();
+    ctx.arc(pendulumX1, pendulumY1, pivotRadius, 0, 2 * Math.PI);
+    ctx.fillStyle = 'black';
+    ctx.fill();
+    ctx.stroke();
+
+    ctx.beginPath();
+    ctx.arc(pendulumX2, pendulumY2, bobRadius, 0, 2 * Math.PI);
+    ctx.fillStyle = '#7DB9B6';
+    ctx.fill();
+    ctx.stroke();
+
+    // Draw torque arc
+    const torqueMagnitude = -action[0];
+    const arrowRadius = canvasWidth * 0.4;
+    const magnitudeRadians = (Math.PI * 2 / 3 * torqueMagnitude);
+    const startAngle =-Math.PI/2 + (torqueMagnitude > 0 ? 0 : magnitudeRadians) - state.theta1;
+    const endAngle = -Math.PI/2 + (torqueMagnitude < 0 ? 0 : magnitudeRadians) - state.theta1;
+
+    ctx.beginPath();
+    ctx.arc(pendulumX1, pendulumY1, arrowRadius, startAngle, endAngle);
+    ctx.strokeStyle = 'black';
+    ctx.lineWidth = canvasWidth * 0.04;
+    ctx.stroke();
+
+    // Draw arrowhead
+    const arrowAngle = torqueMagnitude > 0 ? endAngle : startAngle;
+    const arrowHeadAngularOffset = torqueMagnitude * Math.PI / 180 * 20;
+    const arrowX = pendulumX1 + arrowRadius * Math.cos(arrowAngle + arrowHeadAngularOffset);
+    const arrowY = pendulumY1 + arrowRadius * Math.sin(arrowAngle + arrowHeadAngularOffset);
+
+    const headlen = canvasWidth * 0.20 * Math.min(Math.abs(torqueMagnitude) * 2, 1);
+    const angleOffset = Math.PI / 6;
+    const rotationAngle = Math.PI / 2 + (torqueMagnitude > 0 ? 0 : Math.PI);
+
+    ctx.beginPath();
+    ctx.moveTo(arrowX, arrowY);
+    ctx.lineTo(
+        arrowX - headlen * Math.cos(arrowAngle + arrowHeadAngularOffset / 2 - angleOffset + rotationAngle),
+        arrowY - headlen * Math.sin(arrowAngle + arrowHeadAngularOffset / 2 - angleOffset + rotationAngle)
+    );
+    ctx.lineTo(
+        arrowX - headlen * Math.cos(arrowAngle + arrowHeadAngularOffset / 2 + angleOffset + rotationAngle),
+        arrowY - headlen * Math.sin(arrowAngle + arrowHeadAngularOffset / 2 + angleOffset + rotationAngle)
+    );
+    ctx.lineTo(arrowX, arrowY);
+    ctx.fillStyle = 'black';
+    ctx.fill();
+}
+        )RL_TOOLS_LITERAL";
+        return ui;
+    }
+    template <typename DEVICE, typename SPEC>
+    std::string get_ui(DEVICE& device, rl::environments::AcrobotSwingup<SPEC>& env){
+        return get_ui(device, static_cast<rl::environments::Acrobot<SPEC>&>(env));
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/environments/acrobot/operations_generic.h

@@ -0,0 +1,202 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ENVIRONMENTS_ACROBOT_OPERATIONS_GENERIC_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ENVIRONMENTS_ACROBOT_OPERATIONS_GENERIC_H
+#include "acrobot.h"
+#include "../operations_generic.h"
+// adapted from (and tested agains) https://github.com/Farama-Foundation/Gymnasium/blob/v0.28.1/gymnasium/envs/classic_control/acrobot.py
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::rl::environments::acrobot {
+    template <typename T>
+    RL_TOOLS_FUNCTION_PLACEMENT T clip(T x, T min, T max){
+        x = x < min ? min : (x > max ? max : x);
+        return x;
+    }
+    template <typename DEVICE, typename T>
+    RL_TOOLS_FUNCTION_PLACEMENT T f_mod_python(const DEVICE& dev, T a, T b){
+        return a - b * math::floor(dev, a / b);
+    }
+
+    template <typename DEVICE, typename T>
+    RL_TOOLS_FUNCTION_PLACEMENT T angle_normalize(const DEVICE& dev, T x){
+        return f_mod_python(dev, (x + math::PI<T>), (2 * math::PI<T>)) - math::PI<T>;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, const rl::environments::Acrobot<SPEC>& env){ }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, const rl::environments::Acrobot<SPEC>& env){ }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void init(DEVICE& device, const rl::environments::Acrobot<SPEC>& env){ }
+    template<typename DEVICE, typename SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT void sample_initial_parameters(DEVICE& device, const rl::environments::Acrobot<SPEC>& env, typename rl::environments::Acrobot<SPEC>::Parameters& parameters, RNG& rng){ }
+    template<typename DEVICE, typename SPEC>
+    static void initial_parameters(DEVICE& device, const rl::environments::Acrobot<SPEC>& env, typename rl::environments::Acrobot<SPEC>::Parameters& parameters){ }
+    template<typename DEVICE, typename SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT void sample_initial_state(DEVICE& device, const rl::environments::Acrobot<SPEC>& env, typename rl::environments::Acrobot<SPEC>::Parameters& parameters, typename rl::environments::Acrobot<SPEC>::State& state, RNG& rng){
+        state.theta_1     = random::uniform_real_distribution(typename DEVICE::SPEC::RANDOM(), -0.1, 0.1, rng);
+        state.theta_2     = random::uniform_real_distribution(typename DEVICE::SPEC::RANDOM(), -0.1, 0.1, rng);
+        state.theta_1_dot = random::uniform_real_distribution(typename DEVICE::SPEC::RANDOM(), -0.1, 0.1, rng);
+        state.theta_2_dot = random::uniform_real_distribution(typename DEVICE::SPEC::RANDOM(), -0.1, 0.1, rng);
+    }
+    template<typename DEVICE, typename SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT void sample_initial_state(DEVICE& device, const rl::environments::AcrobotSwingup<SPEC>& env, typename rl::environments::Acrobot<SPEC>::Parameters& parameters, typename rl::environments::AcrobotSwingup<SPEC>::State& state, RNG& rng){
+        using T = typename SPEC::T;
+        state.theta_1     = random::uniform_real_distribution(typename DEVICE::SPEC::RANDOM(), -math::PI<T>, math::PI<T>, rng);
+        state.theta_2     = random::uniform_real_distribution(typename DEVICE::SPEC::RANDOM(), -math::PI<T>, math::PI<T>, rng);
+        state.theta_1_dot = 0;
+        state.theta_2_dot = 0;
+    }
+    template<typename DEVICE, typename SPEC>
+    static void initial_state(DEVICE& device, const rl::environments::Acrobot<SPEC>& env, typename rl::environments::Acrobot<SPEC>::Parameters& parameters, typename rl::environments::Acrobot<SPEC>::State& state){
+        state.theta_1     = 0;
+        state.theta_2     = 0;
+        state.theta_1_dot = 0;
+        state.theta_2_dot = 0;
+    }
+    namespace rl::environments::acrobot{
+        template <typename T, typename PARAMS>
+        RL_TOOLS_FUNCTION_PLACEMENT void dsdt(T state[4], T action, T d_state[4], const PARAMS& params){
+
+            T m1 = params.LINK_MASS_1;
+            T m2 = params.LINK_MASS_2;
+            T l1 = params.LINK_LENGTH_1;
+            T lc1 = params.LINK_COM_POS_1;
+            T lc2 = params.LINK_COM_POS_2;
+            T I1 = params.LINK_MOI;
+            T I2 = params.LINK_MOI;
+            T g = 9.8;
+            T theta1 = state[0];
+            T theta2 = state[1];
+            T dtheta1 = state[2];
+            T dtheta2 = state[3];
+            T d1 = (
+                    m1 * lc1 * lc1
+                    + m2 * (l1*l1 + lc2*lc2 + 2 * l1 * lc2 * cos(theta2))
+                    + I1
+                    + I2
+            );
+            T d2 = m2 * (lc2*lc2 + l1 * lc2 * cos(theta2)) + I2;
+            T phi2 = m2 * lc2 * g * cos(theta1 + theta2 - math::PI<T> / 2.0);
+            T phi1 = (
+                    -m2 * l1 * lc2 * dtheta2 * dtheta2 * sin(theta2)
+                    - 2 * m2 * l1 * lc2 * dtheta2 * dtheta1 * sin(theta2)
+                    + (m1 * lc1 + m2 * l1) * g * cos(theta1 - math::PI<T> / 2)
+                    + phi2
+            );
+
+            T ddtheta2 = (
+                               action + d2 / d1 * phi1 - m2 * l1 * lc2 * dtheta1*dtheta1 * sin(theta2) - phi2
+                       ) / (m2 * lc2*lc2 + I2 - d2*d2 / d1);
+            T ddtheta1 = -(d2 * ddtheta2 + phi1) / d1;
+            d_state[0] = dtheta1;
+            d_state[1] = dtheta2;
+            d_state[2] = ddtheta1;
+            d_state[3] = ddtheta2;
+        }
+
+        template <typename T, typename PARAMS>
+        RL_TOOLS_FUNCTION_PLACEMENT void rk4(T state[4], T action, T next_state[4], T dt, const PARAMS& params){
+
+            T k1[4], k2[4], k3[4], k4[4], y1[4], y2[4], y3[4];
+
+            T dt2 = dt / 2.0;
+
+            rl::environments::acrobot::dsdt(state, action, k1, params);
+            for (int i = 0; i < 4; ++i){
+                y1[i] = state[i] + dt2 * k1[i];
+            }
+            rl::environments::acrobot::dsdt(y1, action, k2, params);
+            for (int i = 0; i < 4; ++i){
+                y2[i] = state[i] + dt2 * k2[i];
+            }
+            rl::environments::acrobot::dsdt(y2, action, k3, params);
+            for (int i = 0; i < 4; ++i){
+                y3[i] = state[i] + dt * k3[i];
+            }
+            rl::environments::acrobot::dsdt(y3, action, k4, params);
+            for (int i = 0; i < 4; ++i){
+                next_state[i] = state[i] + dt / 6.0 * (k1[i] + 2 * k2[i] + 2 * k3[i] + k4[i]);
+            }
+        }
+    }
+    template<typename DEVICE, typename SPEC, typename ACTION_SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT typename SPEC::T step(DEVICE& device, const rl::environments::Acrobot<SPEC>& env, typename rl::environments::Acrobot<SPEC>::Parameters& parameters, const typename rl::environments::Acrobot<SPEC>::State& state, const Matrix<ACTION_SPEC>& action, typename rl::environments::Acrobot<SPEC>::State& next_state, RNG& rng) {
+        static_assert(ACTION_SPEC::ROWS == 1);
+        static_assert(ACTION_SPEC::COLS == 1);
+        using namespace rl::environments::acrobot;
+        using T = typename SPEC::T;
+        using PARAMS = typename SPEC::PARAMETERS;
+
+        T state_flat[4] = {state.theta_1, state.theta_2, state.theta_1_dot, state.theta_2_dot};
+        T next_state_flat[4];
+        T action_clamped = math::clamp(device.math, get(action, 0, 0), (T)-1, (T)1);
+//        action_clamped = 0;
+        T action_scaled = (action_clamped + 1.0) / 2.0 * (PARAMS::MAX_TORQUE - (PARAMS::MIN_TORQUE)) + (PARAMS::MIN_TORQUE);
+        rl::environments::acrobot::rk4(state_flat, action_scaled, next_state_flat, PARAMS::DT, PARAMS{});
+
+        next_state_flat[0] = angle_normalize(device.math, next_state_flat[0]);
+        next_state_flat[1] = angle_normalize(device.math, next_state_flat[1]);
+        next_state_flat[2] = math::clamp(    device.math, next_state_flat[2], -PARAMS::MAX_VEL_1, PARAMS::MAX_VEL_1);
+        next_state_flat[3] = math::clamp(    device.math, next_state_flat[3], -PARAMS::MAX_VEL_2, PARAMS::MAX_VEL_2);
+
+        next_state.theta_1 = next_state_flat[0];
+        next_state.theta_2 = next_state_flat[1];
+        next_state.theta_1_dot = next_state_flat[2];
+        next_state.theta_2_dot = next_state_flat[3];
+
+        return SPEC::PARAMETERS::DT;
+    }
+    template<typename DEVICE, typename SPEC, typename ACTION_SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT static typename SPEC::T reward(DEVICE& device, const rl::environments::Acrobot<SPEC>& env, const typename rl::environments::Acrobot<SPEC>::Parameters& parameters, const typename rl::environments::Acrobot<SPEC>::State& state, const Matrix<ACTION_SPEC>& action, const typename rl::environments::Acrobot<SPEC>::State& next_state, RNG& rng){
+        return -1;
+    }
+    template<typename DEVICE, typename SPEC, typename ACTION_SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT static typename SPEC::T reward(DEVICE& device, const rl::environments::AcrobotSwingup<SPEC>& env, const typename rl::environments::Acrobot<SPEC>::Parameters& parameters, const typename rl::environments::Acrobot<SPEC>::State& state, const Matrix<ACTION_SPEC>& action, const typename rl::environments::Acrobot<SPEC>::State& next_state, RNG& rng){
+        using T = typename SPEC::T;
+//        return (-cos(state.theta_1) * SPEC::PARAMETERS::LINK_LENGTH_1 - cos(state.theta_2 + state.theta_1) * SPEC::PARAMETERS::LINK_LENGTH_2) * 0.01;
+        typename DEVICE::SPEC::MATH dm;
+        T a_angle_cost = rl::environments::acrobot::angle_normalize(dm, next_state.theta_1 - math::PI<T>);
+//        a_angle_cost *= a_angle_cost;
+        a_angle_cost = math::abs(device.math, a_angle_cost);
+        T b_angle_cost = rl::environments::acrobot::angle_normalize(dm, next_state.theta_2);
+//        b_angle_cost *= b_angle_cost;
+        b_angle_cost = math::abs(device.math, b_angle_cost);
+
+        T a_vel_cost = next_state.theta_1_dot;
+        a_vel_cost *= a_vel_cost;
+        T b_vel_cost = next_state.theta_2_dot;
+        b_vel_cost *= b_vel_cost;
+        T torque_cost = get(action, 0, 0);
+        torque_cost *= torque_cost;
+
+//        return math::exp(dm, -a_angle_cost) + math::exp(dm, -b_angle_cost) + SPEC::PARAMETERS::VEL_PENALTY * math::exp(dm, -a_vel_cost) + SPEC::PARAMETERS::VEL_PENALTY*math::exp(dm, -b_vel_cost) + 0.3 * math::exp(dm, -torque_cost);
+        return (-a_angle_cost -b_angle_cost)/100*SPEC::PARAMETERS::DT/0.02;
+    }
+
+    template<typename DEVICE, typename SPEC, typename OBS_TYPE_SPEC, typename OBS_SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT static void observe(DEVICE& device, const rl::environments::Acrobot<SPEC>& env, const typename rl::environments::Acrobot<SPEC>::Parameters& parameters, const typename rl::environments::Acrobot<SPEC>::State& state, const rl::environments::acrobot::Observation<OBS_TYPE_SPEC>&, Matrix<OBS_SPEC>& observation, RNG& rng){
+        static_assert(OBS_SPEC::ROWS == 1);
+        static_assert(OBS_SPEC::COLS == 6);
+        typedef typename SPEC::T T;
+        set(observation, 0, 0, math::cos(device.math, state.theta_1));
+        set(observation, 0, 1, math::sin(device.math, state.theta_1));
+        set(observation, 0, 2, math::cos(device.math, state.theta_2));
+        set(observation, 0, 3, math::sin(device.math, state.theta_2));
+        set(observation, 0, 4, state.theta_1_dot);
+        set(observation, 0, 5, state.theta_2_dot);
+    }
+    template<typename DEVICE, typename SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT static bool terminated(DEVICE& device, const rl::environments::Acrobot<SPEC>& env, const typename rl::environments::Acrobot<SPEC>::Parameters& parameters, const typename rl::environments::Acrobot<SPEC>::State state, RNG& rng){
+        return (-cos(state.theta_1) - cos(state.theta_2 + state.theta_1)) > 1.0;
+    }
+    template<typename DEVICE, typename SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT static bool terminated(DEVICE& device, const rl::environments::AcrobotSwingup<SPEC>& env, const typename rl::environments::Acrobot<SPEC>::Parameters& parameters, const typename rl::environments::Acrobot<SPEC>::State state, RNG& rng){
+        return false; //math::abs(typename DEVICE::SPEC::MATH{}, state.theta_1_dot) > SPEC::PARAMETERS::MAX_VEL_1 || math::abs(typename DEVICE::SPEC::MATH{}, state.theta_1_dot) > SPEC::PARAMETERS::MAX_VEL_2;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif