hyperlux 0.0.1__py3-none-any.whl

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/environments/mujoco/ant/persist.h

@@ -0,0 +1,86 @@
+#include "../../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ENVIRONMENTS_MUJOCO_ANT_PERSIST_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ENVIRONMENTS_MUJOCO_ANT_PERSIST_H
+#include "ant.h"
+#include <cstring>
+#include <vector>
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template <typename DEVICE, typename SPEC, typename GROUP>
+    void save(DEVICE& device, rl::environments::mujoco::Ant<SPEC>& environment, GROUP& group){
+        auto* model = environment.model;
+        auto* data = environment.data;
+        utils::assert_exit(device, model->nplugin == 0, "Ant persistence requires a model without simulator plugins");
+        const int model_size = mj_sizeModel(model);
+        std::vector<char> model_data(model_size);
+        mj_saveModel(model, nullptr, model_data.data(), model_size);
+        save_binary(device, &model_size, 1, group, "model_size");
+        save_binary(device, model_data.data(), model_size, group, "model");
+        save_binary(device, environment.init_q, SPEC::STATE_DIM_Q, group, "initial_position");
+        save_binary(device, environment.init_q_dot, SPEC::STATE_DIM_Q_DOT, group, "initial_velocity");
+        save_binary(device, &environment.last_reward, 1, group, "last_reward");
+        save_binary(device, &environment.last_terminated, 1, group, "last_terminated");
+        save_binary(device, &data->time, 1, group, "time");
+        save_binary(device, data->qpos, model->nq * 1, group, "qpos");
+        save_binary(device, data->qvel, model->nv * 1, group, "qvel");
+        save_binary(device, data->act, model->na * 1, group, "act");
+        save_binary(device, data->qacc_warmstart, model->nv * 1, group, "qacc_warmstart");
+        save_binary(device, data->ctrl, model->nu * 1, group, "ctrl");
+        save_binary(device, data->qfrc_applied, model->nv * 1, group, "qfrc_applied");
+        save_binary(device, data->xfrc_applied, model->nbody * 6, group, "xfrc_applied");
+        save_binary(device, data->mocap_pos, model->nmocap * 3, group, "mocap_pos");
+        save_binary(device, data->mocap_quat, model->nmocap * 4, group, "mocap_quat");
+        save_binary(device, data->userdata, model->nuserdata * 1, group, "userdata");
+        save_binary(device, data->xpos, model->nbody * 3, group, "xpos");
+    }
+    template <typename DEVICE, typename SPEC, typename GROUP>
+    bool load(DEVICE& device, rl::environments::mujoco::Ant<SPEC>& environment, GROUP& group){
+        int model_size = 0;
+        if(!load_binary(device, &model_size, 1, group, "model_size") || model_size <= 0) return false;
+        std::vector<char> model_data(model_size);
+        if(!load_binary(device, model_data.data(), model_size, group, "model")) return false;
+        mjVFS vfs;
+        mj_defaultVFS(&vfs);
+        if(mj_makeEmptyFileVFS(&vfs, "checkpoint.mjb", model_size) != 0){
+            mj_deleteVFS(&vfs);
+            return false;
+        }
+        std::memcpy(vfs.filedata[mj_findFileVFS(&vfs, "checkpoint.mjb")], model_data.data(), model_size);
+        auto* model = mj_loadModel("checkpoint.mjb", &vfs);
+        mj_deleteVFS(&vfs);
+        if(model == nullptr) return false;
+        if(model->nq != SPEC::STATE_DIM_Q || model->nv != SPEC::STATE_DIM_Q_DOT || model->nu != SPEC::ACTION_DIM || model->nplugin != 0){
+            mj_deleteModel(model);
+            return false;
+        }
+        auto* data = mj_makeData(model);
+        if(data == nullptr){ mj_deleteModel(model); return false; }
+        mj_deleteData(environment.data);
+        mj_deleteModel(environment.model);
+        environment.model = model;
+        environment.data = data;
+        environment.torso_id = mj_name2id(model, mjOBJ_XBODY, "torso");
+        bool success = load_binary(device, environment.init_q, SPEC::STATE_DIM_Q, group, "initial_position");
+        success &= load_binary(device, environment.init_q_dot, SPEC::STATE_DIM_Q_DOT, group, "initial_velocity");
+        success &= load_binary(device, &environment.last_reward, 1, group, "last_reward");
+        success &= load_binary(device, &environment.last_terminated, 1, group, "last_terminated");
+        success &= load_binary(device, &data->time, 1, group, "time");
+        success &= load_binary(device, data->qpos, model->nq * 1, group, "qpos");
+        success &= load_binary(device, data->qvel, model->nv * 1, group, "qvel");
+        success &= load_binary(device, data->act, model->na * 1, group, "act");
+        success &= load_binary(device, data->ctrl, model->nu * 1, group, "ctrl");
+        success &= load_binary(device, data->qfrc_applied, model->nv * 1, group, "qfrc_applied");
+        success &= load_binary(device, data->xfrc_applied, model->nbody * 6, group, "xfrc_applied");
+        success &= load_binary(device, data->mocap_pos, model->nmocap * 3, group, "mocap_pos");
+        success &= load_binary(device, data->mocap_quat, model->nmocap * 4, group, "mocap_quat");
+        success &= load_binary(device, data->userdata, model->nuserdata * 1, group, "userdata");
+        mj_forward(model, data);
+        success &= load_binary(device, data->qacc_warmstart, model->nv, group, "qacc_warmstart");
+        success &= load_binary(device, data->xpos, model->nbody * 3, group, "xpos");
+        return success;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/environments/mujoco/ant/ui.h

@@ -0,0 +1,141 @@
+#include "../../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ENVIRONMENTS_MUJOCO_ANT_UI_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ENVIRONMENTS_MUJOCO_ANT_UI_H
+#include <GLFW/glfw3.h>
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::rl::environments::mujoco::ant {
+    template <typename T_ENVIRONMENT>
+    struct UI{
+        using ENVIRONMENT = T_ENVIRONMENT;
+        ENVIRONMENT* env;
+        GLFWwindow* window;
+        mjvCamera camera;
+        mjvOption option;
+        mjvScene scene;
+        mjrContext context;
+        bool button_left = false;
+        bool button_middle = false;
+        bool button_right =  false;
+        double lastx = 0;
+        double lasty = 0;
+    };
+    namespace ui::callbacks{
+        template <typename UI>
+        void keyboard(GLFWwindow* window, int key, int scancode, int act, int mods) {
+            UI* ui = (UI*)glfwGetWindowUserPointer(window);
+            if (act==GLFW_PRESS && key==GLFW_KEY_BACKSPACE) {
+                mj_resetData(ui->env->model, ui->env->data);
+                mj_forward(ui->env->model, ui->env->data);
+            }
+        }
+        template <typename UI>
+        void mouse_button(GLFWwindow* window, int button, int act, int mods) {
+            UI* ui = (UI*)glfwGetWindowUserPointer(window);
+            ui->button_left = (glfwGetMouseButton(window, GLFW_MOUSE_BUTTON_LEFT)==GLFW_PRESS);
+            ui->button_middle = (glfwGetMouseButton(window, GLFW_MOUSE_BUTTON_MIDDLE)==GLFW_PRESS);
+            ui->button_right = (glfwGetMouseButton(window, GLFW_MOUSE_BUTTON_RIGHT)==GLFW_PRESS);
+
+            glfwGetCursorPos(window, &ui->lastx, &ui->lasty);
+        }
+
+        template <typename UI>
+        void mouse_move(GLFWwindow* window, double xpos, double ypos) {
+            UI* ui = (UI*)glfwGetWindowUserPointer(window);
+            if (!ui->button_left && !ui->button_middle && !ui->button_right) {
+                return;
+            }
+            double dx = xpos - ui->lastx;
+            double dy = ypos - ui->lasty;
+            ui->lastx = xpos;
+            ui->lasty = ypos;
+            int width, height;
+            glfwGetWindowSize(window, &width, &height);
+            bool mod_shift = (glfwGetKey(window, GLFW_KEY_LEFT_SHIFT)==GLFW_PRESS || glfwGetKey(window, GLFW_KEY_RIGHT_SHIFT)==GLFW_PRESS);
+            mjtMouse action;
+            if (ui->button_right) {
+                action = mod_shift ? mjMOUSE_MOVE_H : mjMOUSE_MOVE_V;
+            } else if (ui->button_left) {
+                action = mod_shift ? mjMOUSE_ROTATE_H : mjMOUSE_ROTATE_V;
+            } else {
+                action = mjMOUSE_ZOOM;
+            }
+            mjv_moveCamera(ui->env->model, action, dx/height, dy/height, &ui->scene, &ui->camera);
+        }
+
+        template <typename UI>
+        void scroll(GLFWwindow* window, double xoffset, double yoffset) {
+            UI* ui = (UI*)glfwGetWindowUserPointer(window);
+            mjv_moveCamera(ui->env->model, mjMOUSE_ZOOM, 0, -0.05*yoffset, &ui->scene, &ui->camera);
+        }
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template <typename DEVICE, typename ENVIRONMENT>
+    void init(DEVICE& dev, ENVIRONMENT& env, typename ENVIRONMENT::Parameters& parameters, rl::environments::mujoco::ant::UI<ENVIRONMENT>& ui){
+        using UI = rl::environments::mujoco::ant::UI<ENVIRONMENT>;
+        ui.env = &env;
+        if (!glfwInit()) {
+            mju_error("Could not initialize GLFW");
+        }
+        ui.window = glfwCreateWindow(1200, 900, "Demo", NULL, NULL);
+        glfwSetWindowUserPointer(ui.window, &ui);
+        glfwMakeContextCurrent(ui.window);
+        glfwSwapInterval(1);
+
+        mjv_defaultCamera(&ui.camera);
+        ui.camera.type = mjCAMERA_TRACKING;
+        ui.camera.trackbodyid = env.torso_id;
+        ui.camera.distance = 6;
+        mjv_defaultOption(&ui.option);
+        mjv_defaultScene(&ui.scene);
+        mjr_defaultContext(&ui.context);
+
+        mjv_makeScene(ui.env->model, &ui.scene, 2000);
+        mjr_makeContext(ui.env->model, &ui.context, mjFONTSCALE_150);
+
+        glfwSetKeyCallback(ui.window, rl::environments::mujoco::ant::ui::callbacks::keyboard<UI>);
+        glfwSetCursorPosCallback(ui.window, rl::environments::mujoco::ant::ui::callbacks::mouse_move<UI>);
+        glfwSetMouseButtonCallback(ui.window, rl::environments::mujoco::ant::ui::callbacks::mouse_button<UI>);
+        glfwSetScrollCallback(ui.window, rl::environments::mujoco::ant::ui::callbacks::scroll<UI>);
+    }
+    template <typename DEVICE, typename ENVIRONMENT>
+    void set_state(DEVICE& dev, ENVIRONMENT& environment, typename ENVIRONMENT::Parameters& parameters, rl::environments::mujoco::ant::UI<ENVIRONMENT>& ui, const typename ENVIRONMENT::State& state){
+        using TI = typename DEVICE::index_t;
+        for(TI state_i = 0; state_i < ENVIRONMENT::SPEC::STATE_DIM_Q; state_i++){
+            ui.env->data->qpos[state_i] = state.q[state_i];
+        }
+        for(TI state_i = 0; state_i < ENVIRONMENT::SPEC::STATE_DIM_Q_DOT; state_i++){
+            ui.env->data->qvel[state_i] = state.q_dot[state_i];
+        }
+        mj_forward(ui.env->model, ui.env->data);
+        mjrRect viewport = {0, 0, 0, 0};
+        glfwGetFramebufferSize(ui.window, &viewport.width, &viewport.height);
+
+        mjv_updateScene(ui.env->model, ui.env->data, &ui.option, NULL, &ui.camera, mjCAT_ALL, &ui.scene);
+        mjr_render(viewport, &ui.scene, &ui.context);
+
+        glfwSwapBuffers(ui.window);
+
+        glfwPollEvents();
+    }
+
+    template <typename DEVICE, typename ENVIRONMENT>
+    void free(DEVICE& dev, ENVIRONMENT& env, rl::environments::mujoco::ant::UI<ENVIRONMENT>& ui){
+        mjv_freeScene(&ui.scene);
+        mjr_freeContext(&ui.context);
+
+#if defined(__APPLE__) || defined(_WIN32)
+        glfwTerminate();
+#endif
+    }
+
+
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif

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

@@ -0,0 +1 @@
+#include "ant/ant.h"

hyperdrone/_vendor/rl-tools/include/rl_tools/rl/environments/multi_agent/bottleneck/bottleneck.h

@@ -0,0 +1,113 @@
+#include "../../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ENVIRONMENTS_MULTI_AGENT_BOTTLENECK_BOTTLENECK_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ENVIRONMENTS_MULTI_AGENT_BOTTLENECK_BOTTLENECK_H
+
+#include "../../../../math/operations_generic.h"
+#include "../environments.h"
+#include "../../observation.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::rl::environments::multi_agent::bottleneck {
+    template <typename T_T, typename T_TI>
+    struct DefaultParameters {
+        using T = T_T;
+        using TI = T_TI;
+        static constexpr TI N_AGENTS = 2;
+        static constexpr TI LIDAR_RESOLUTION = 5;
+        static constexpr T LIDAR_FOV = math::PI<T> * 20/180; // in radians (0 to PI)
+        static constexpr T LIDAR_RANGE = 3;
+        static constexpr T DT = 0.05;
+        static constexpr T ARENA_WIDTH = 20; // x: left to right, y: top to bottom (display coordinate system)
+        static constexpr T ARENA_HEIGHT = 20;
+        static constexpr T AGENT_DIAMETER = 1.5;
+        static constexpr T AGENT_MAX_SPEED = 2;
+        static constexpr T AGENT_MAX_ACCELERATION = 2;
+        static constexpr T AGENT_MAX_ANGULAR_VELOCITY = 2;
+        static constexpr T AGENT_MAX_ANGULAR_ACCELERATION = 2;
+        static constexpr T BOTTLENECK_POSITION = 10; // y
+        static constexpr T BOTTLENECK_WIDTH = 2;
+        static constexpr T BARRIER_WIDTH = 0.5;
+        static constexpr TI EPISODE_STEP_LIMIT = 200;
+        static constexpr bool SPAWN_BOTH_SIDES = false;
+    };
+    template <typename T_PARAMETERS>
+    struct Observation{
+        using PARAMETERS = T_PARAMETERS;
+        using T = typename PARAMETERS::T;
+        using TI = typename PARAMETERS::TI;
+        static constexpr TI PER_AGENT_DIM = 7 + PARAMETERS::LIDAR_RESOLUTION;
+        static constexpr TI DIM = PARAMETERS::N_AGENTS * PER_AGENT_DIM;
+        using SHAPE = tensor::Shape<TI, DIM>;
+    };
+    template <typename T_PARAMETERS>
+    struct ObservationPrivileged{
+        using PARAMETERS = T_PARAMETERS;
+        using T = typename PARAMETERS::T;
+        using TI = typename PARAMETERS::TI;
+        static constexpr TI PER_AGENT_DIM = 7 + PARAMETERS::LIDAR_RESOLUTION;
+        static constexpr TI DIM = PARAMETERS::N_AGENTS * PER_AGENT_DIM;
+        using SHAPE = tensor::Shape<TI, DIM>;
+    };
+    template <typename T_TYPE_POLICY, typename T_TI, typename T_PARAMETERS = DefaultParameters<typename T_TYPE_POLICY::DEFAULT, T_TI>, typename T_OBSERVATION = Observation<T_PARAMETERS>, typename T_OBSERVATION_PRIVILEGED = ObservationPrivileged<T_PARAMETERS>>
+    struct Specification{
+        using T = typename T_TYPE_POLICY::DEFAULT;
+        using TI = T_TI;
+        using OBSERVATION = T_OBSERVATION;
+        using OBSERVATION_PRIVILEGED = T_OBSERVATION_PRIVILEGED;
+        using PARAMETERS = T_PARAMETERS;
+    };
+
+    template <typename T>
+    struct Intersection {
+        bool intersects = false;
+        T point[2];
+        T distance;
+    };
+
+    template <typename T, typename TI>
+    struct AgentState {
+        T position[2];
+        T orientation;
+        T velocity[2];
+        T angular_velocity;
+        Intersection<T> lidar[DefaultParameters<T, TI>::LIDAR_RESOLUTION];
+        bool dead;
+    };
+
+    template <typename SPEC>
+    struct State{
+        using T = typename SPEC::T;
+        using TI = typename SPEC::TI;
+        AgentState<T, TI> agent_states[SPEC::PARAMETERS::N_AGENTS];
+    };
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::rl::environments::multi_agent{
+    template <typename T_SPEC>
+    struct Bottleneck: Environment<typename T_SPEC::T, typename T_SPEC::TI>{
+        using SPEC = T_SPEC;
+        using PARAMETERS = typename SPEC::PARAMETERS;
+        using T = typename SPEC::T;
+        using TI = typename SPEC::TI;
+        using State = multi_agent::bottleneck::State<SPEC>;
+        using Parameters = typename SPEC::PARAMETERS;
+        using Observation = typename SPEC::OBSERVATION;
+        using ObservationPrivileged = typename SPEC::OBSERVATION_PRIVILEGED;
+        static constexpr TI N_AGENTS = Parameters::N_AGENTS;
+        static constexpr TI PER_AGENT_ACTION_DIM = 2; // linear and angular acceleration
+        static constexpr TI ACTION_DIM = N_AGENTS * PER_AGENT_ACTION_DIM;
+        static constexpr TI EPISODE_STEP_LIMIT = Parameters::EPISODE_STEP_LIMIT;
+    };
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+
+
+
+
+
+
+#endif

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

@@ -0,0 +1,289 @@
+#include "../../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_RL_ENVIRONMENTS_MULTI_AGENT_BOTTLENECK_OPERATIONS_CPU_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_RL_ENVIRONMENTS_MULTI_AGENT_BOTTLENECK_OPERATIONS_CPU_H
+
+#include "bottleneck.h"
+#include "operations_generic.h"
+
+#include <string>
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template <typename DEVICE, typename SPEC>
+    std::string json(DEVICE&, rl::environments::multi_agent::Bottleneck<SPEC>& env){
+        std::string json = "{";
+        json += "}";
+        return json;
+    }
+    template <typename DEVICE, typename SPEC>
+    std::string json(DEVICE&, rl::environments::multi_agent::Bottleneck<SPEC>& env, const typename rl::environments::multi_agent::Bottleneck<SPEC>::Parameters& parameters){
+        std::string json = "{";
+        json += "\"N_AGENTS\":" + std::to_string(SPEC::PARAMETERS::N_AGENTS) + ",";
+        json += "\"LIDAR_RESOLUTION\":" + std::to_string(SPEC::PARAMETERS::LIDAR_RESOLUTION) + ",";
+        json += "\"LIDAR_FOV\":" + std::to_string(SPEC::PARAMETERS::LIDAR_FOV) + ",";
+        json += "\"LIDAR_RANGE\":" + std::to_string(SPEC::PARAMETERS::LIDAR_RANGE) + ",";
+        json += "\"DT\":" + std::to_string(SPEC::PARAMETERS::DT) + ",";
+        json += "\"ARENA_WIDTH\":" + std::to_string(SPEC::PARAMETERS::ARENA_WIDTH) + ",";
+        json += "\"ARENA_HEIGHT\":" + std::to_string(SPEC::PARAMETERS::ARENA_HEIGHT) + ",";
+        json += "\"AGENT_DIAMETER\":" + std::to_string(SPEC::PARAMETERS::AGENT_DIAMETER) + ",";
+        json += "\"AGENT_MAX_SPEED\":" + std::to_string(SPEC::PARAMETERS::AGENT_MAX_SPEED) + ",";
+        json += "\"AGENT_MAX_ACCELERATION\":" + std::to_string(SPEC::PARAMETERS::AGENT_MAX_ACCELERATION) + ",";
+        json += "\"AGENT_MAX_ANGULAR_VELOCITY\":" + std::to_string(SPEC::PARAMETERS::AGENT_MAX_ANGULAR_VELOCITY) + ",";
+        json += "\"AGENT_MAX_ANGULAR_ACCELERATION\":" + std::to_string(SPEC::PARAMETERS::AGENT_MAX_ANGULAR_ACCELERATION) + ",";
+        json += "\"BOTTLENECK_POSITION\":" + std::to_string(SPEC::PARAMETERS::BOTTLENECK_POSITION) + ",";
+        json += "\"BOTTLENECK_WIDTH\":" + std::to_string(SPEC::PARAMETERS::BOTTLENECK_WIDTH) + ",";
+        json += "\"BARRIER_WIDTH\":" + std::to_string(SPEC::PARAMETERS::BARRIER_WIDTH);
+        json += "}";
+        return json;
+    }
+    template <typename DEVICE, typename SPEC>
+    std::string json(DEVICE&, rl::environments::multi_agent::Bottleneck<SPEC>& env, const typename rl::environments::multi_agent::Bottleneck<SPEC>::Parameters& parameters, const typename rl::environments::multi_agent::Bottleneck<SPEC>::State& state){
+        using TI = typename DEVICE::index_t;
+        std::string agent_states = "[";
+        for(TI agent_i = 0; agent_i < SPEC::PARAMETERS::N_AGENTS; agent_i++){
+            std::string agent_state = "{";
+            agent_state += "\"position\": [" + std::to_string(state.agent_states[agent_i].position[0]) + "," + std::to_string(state.agent_states[agent_i].position[1]) + "],";
+            agent_state += "\"orientation\": " + std::to_string(state.agent_states[agent_i].orientation) + ",";
+            agent_state += "\"velocity\": [" + std::to_string(state.agent_states[agent_i].velocity[0]) + "," + std::to_string(state.agent_states[agent_i].velocity[1]) + "],";
+            agent_state += "\"angular_velocity\": " + std::to_string(state.agent_states[agent_i].angular_velocity) + ", ";
+            agent_state += "\"dead\": " + std::to_string(state.agent_states[agent_i].dead);
+            std::string lidar = "[";
+            for(TI i = 0; i < SPEC::PARAMETERS::LIDAR_RESOLUTION; i++){
+                std::string ray = "{";
+                ray += "\"distance\": " + std::to_string(state.agent_states[agent_i].lidar[i].distance) + ",";
+                ray += "\"point\": [" + std::to_string(state.agent_states[agent_i].lidar[i].point[0]) + "," + std::to_string(state.agent_states[agent_i].lidar[i].point[1]) + "], ";
+                ray += "\"intersects\": " + std::to_string(state.agent_states[agent_i].lidar[i].intersects);
+                ray += "}";
+                lidar += ray;
+                if(i < SPEC::PARAMETERS::LIDAR_RESOLUTION - 1){
+                    lidar += ",";
+                }
+            }
+            lidar += "]";
+            agent_state += ", \"lidar\": " + lidar;
+            agent_state += "}";
+            agent_states += agent_state;
+            if(agent_i < parameters.N_AGENTS - 1){
+                agent_states += ",";
+            }
+        }
+        agent_states += "]";
+        std::string json = "{";
+        json += "\"agent_states\": " + agent_states;
+        json += "}";
+        return json;
+    }
+
+    template <typename DEVICE, typename SPEC>
+    std::string get_ui(DEVICE& device, rl::environments::multi_agent::Bottleneck<SPEC>& env){
+        // just the body of `function render(ctx, state, action) {` (so that it can be easily processed by `new Function("ctx", "state", "action", body)`
+        std::string ui = R"RL_TOOLS_LITERAL(
+export async function init(canvas, parameters, 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 canvasWidth = ctx.canvas.width;
+    const canvasHeight = ctx.canvas.height;
+    const scaleX = canvasWidth / parameters.ARENA_WIDTH;
+    const scaleY = canvasHeight / parameters.ARENA_HEIGHT;
+    console.assert(scaleX == scaleY);
+
+
+    // Draw large, gray arrow in the background
+    const arrowStartX = canvasWidth / 5;
+    const arrowEndX = 4 * canvasWidth / 5;
+    const arrowY = canvasHeight / 2;
+    const arrowThickness = canvasHeight / 10;
+    const arrowHeadLength = arrowThickness * 2.0;
+    const arrowHeadWidth = arrowThickness * 3;
+    const arrowColor = '#f3f3f3'
+    ctx.beginPath();
+    ctx.moveTo(arrowStartX, arrowY);
+    ctx.lineTo(arrowEndX - arrowHeadLength, arrowY);
+    ctx.strokeStyle = arrowColor;
+    ctx.lineWidth = arrowThickness;
+    ctx.stroke();
+    ctx.beginPath();
+    ctx.moveTo(arrowEndX, arrowY);
+    ctx.lineTo(arrowEndX - arrowHeadLength, arrowY - arrowHeadWidth / 2);
+    ctx.lineTo(arrowEndX - arrowHeadLength, arrowY + arrowHeadWidth / 2);
+    ctx.closePath();
+    ctx.fillStyle = arrowColor;
+    ctx.fill();
+
+
+    // Draw the bottleneck barrier
+    const barrierX = (parameters.ARENA_WIDTH / 2 - parameters.BARRIER_WIDTH / 2) * scaleX;
+    const barrierWidth = parameters.BARRIER_WIDTH * scaleX;
+    const bottleneckTopY = (parameters.BOTTLENECK_POSITION - parameters.BOTTLENECK_WIDTH / 2) * scaleY;
+    const bottleneckBottomY = (parameters.BOTTLENECK_POSITION + parameters.BOTTLENECK_WIDTH / 2) * scaleY;
+    ctx.fillStyle = 'gray';
+    ctx.fillRect(barrierX, 0, barrierWidth, bottleneckTopY);
+    ctx.fillRect(barrierX, bottleneckBottomY, barrierWidth, canvasHeight - bottleneckBottomY);
+
+    const agentRadius = parameters.AGENT_DIAMETER * scaleX / 2;
+
+    // Draw agents and their actions
+    for (let i = 0; i < parameters.N_AGENTS; i++) {
+        const agent = state.agent_states[i];
+        const posX = agent.position[0] * scaleX;
+        const posY = agent.position[1] * scaleY;
+        const orientation = agent.orientation;
+
+        // Draw lidar
+        if(!agent.dead){
+            for (let lidar_i = 0; lidar_i < agent.lidar.length; lidar_i++) {
+                const lidar = agent.lidar[lidar_i];
+                if (lidar.intersects) {
+                    ctx.beginPath();
+                    ctx.moveTo(posX, posY);
+                    const lidarEndX = lidar.point[0] * scaleX;
+                    const lidarEndY = lidar.point[1] * scaleY;
+                    ctx.lineTo(lidarEndX, lidarEndY);
+                    ctx.strokeStyle = 'orange';
+                    ctx.lineWidth = 1/25*scaleX;
+                    ctx.stroke();
+                    ctx.beginPath();
+                    const intersectionDotRadius = 3/25*scaleX;
+                    ctx.arc(lidarEndX, lidarEndY, intersectionDotRadius, 0, 2 * Math.PI);
+                    ctx.fillStyle = 'orange';
+                    ctx.fill();
+                }
+            }
+        }
+
+        // Draw agent body
+        ctx.beginPath();
+        ctx.arc(posX, posY, agentRadius, 0, 2 * Math.PI);
+        const primaryColor = '#7DB9B6';
+        ctx.fillStyle = agent.dead ? 'grey' : primaryColor;
+        ctx.fill();
+
+        // Draw the agent ID
+        const labelPosX = posX - agentRadius / 2 * Math.cos(orientation);
+        const labelPosY = posY - agentRadius / 2 * Math.sin(orientation);
+        ctx.save();
+        ctx.translate(labelPosX, labelPosY);
+        ctx.rotate(orientation + Math.PI / 2);
+        const agentID = i.toString();
+        const fontSize = 0.5 * agentRadius;
+        ctx.font = `${fontSize}px Arial`;
+        const textMetrics = ctx.measureText(agentID);
+        const textWidth = textMetrics.width;
+        const textHeight = fontSize;
+        const circleRadius = Math.max(textWidth, textHeight) * 0.70;
+        ctx.beginPath();
+        ctx.arc(0, 0, circleRadius, 0, 2 * Math.PI);
+        ctx.fillStyle = 'white';
+        ctx.fill();
+        ctx.strokeStyle = 'black';
+        ctx.lineWidth = 0.5/25*scaleX;
+        ctx.stroke();
+        ctx.fillStyle = 'black';
+        ctx.textAlign = 'center';
+        ctx.textBaseline = 'middle';
+        const textYOffset = fontSize * 0.1;
+        ctx.fillText(agentID, 0, textYOffset);
+        ctx.restore();
+
+        // Draw agent orientation
+        const endX = posX + agentRadius * Math.cos(orientation);
+        const endY = posY + agentRadius * Math.sin(orientation);
+        ctx.beginPath();
+        ctx.moveTo(posX, posY);
+        ctx.lineTo(endX, endY);
+        ctx.strokeStyle = 'white';
+        ctx.lineWidth = 2/25*scaleX;
+        ctx.stroke();
+
+        // Draw actions (acceleration vectors)
+        const agent_action = [action[i*2 + 0], action[i*2 + 1]].map(action => Math.max(-1, Math.min(1, action)));
+
+        // Linear acceleration in the direction of orientation
+        const accelerationArrowColor = '#dc143c';
+        if(!agent.dead){
+            const accelMagnitude = agent_action[0]/2 * agentRadius;
+            const accelX = accelMagnitude * Math.cos(orientation);
+            const accelY = accelMagnitude * Math.sin(orientation);
+            const offsetX = posX + agentRadius * 0.5 * Math.cos(orientation);
+            const offsetY = posY + agentRadius * 0.5 * Math.sin(orientation);
+            ctx.beginPath();
+            ctx.moveTo(offsetX, offsetY);
+            ctx.lineTo(offsetX + accelX, offsetY + accelY);
+
+            ctx.strokeStyle = accelerationArrowColor;
+            ctx.lineWidth = 2/25*scaleX;
+            ctx.stroke();
+
+            // Draw arrowhead for linear acceleration
+            const angle = Math.atan2(accelY, accelX);
+            const headlen = 0.5 * Math.min(0.5, Math.abs(agent_action[0])) * scaleX;
+            ctx.beginPath();
+            const prolong = 1.1
+            ctx.moveTo(offsetX + accelX*prolong, offsetY + accelY*prolong);
+            ctx.lineTo(offsetX + accelX*prolong - headlen * Math.cos(angle - Math.PI / 6), offsetY + accelY*prolong - headlen * Math.sin(angle - Math.PI / 6));
+            ctx.moveTo(offsetX + accelX*prolong, offsetY + accelY*prolong);
+            ctx.lineTo(offsetX + accelX*prolong - headlen * Math.cos(angle + Math.PI / 6), offsetY + accelY*prolong - headlen * Math.sin(angle + Math.PI / 6));
+            ctx.lineWidth = 1/25*scaleX;
+            ctx.stroke();
+        }
+
+        // Draw circular arrow for angular acceleration
+        if(!agent.dead){
+            const angularAccel = agent_action[1]; // Negative sign to match the canvas coordinate system
+            const direction = Math.sign(angularAccel);
+            const arrowRadius = agentRadius * 1.5;
+            const arrowAngle = Math.PI / 3;
+            const startAngle = orientation;
+            const endAngle = orientation + arrowAngle * angularAccel;
+            const arrowHeadSize = 10 * Math.abs(angularAccel);
+
+            ctx.beginPath();
+            ctx.arc(posX, posY, arrowRadius, startAngle, endAngle, angularAccel < 0);
+            const torqueArrowColor = '#007bff'
+            ctx.strokeStyle = torqueArrowColor;
+            ctx.lineWidth = 2/25*scaleX;
+            ctx.stroke();
+
+            // Draw arrowhead for angular acceleration
+            const arrowHeadAngle = endAngle - direction * arrowAngle * 0.10 - Math.PI / 2;
+            const arrowHeadX = posX + arrowRadius * Math.cos(endAngle);
+            const arrowHeadY = posY + arrowRadius * Math.sin(endAngle);
+
+            ctx.beginPath();
+            ctx.moveTo(arrowHeadX, arrowHeadY);
+            ctx.lineTo(
+                arrowHeadX + direction * arrowHeadSize * Math.cos(arrowHeadAngle - Math.PI / 6),
+                arrowHeadY + direction * arrowHeadSize * Math.sin(arrowHeadAngle - Math.PI / 6)
+            );
+            ctx.moveTo(arrowHeadX, arrowHeadY);
+            ctx.lineTo(
+                arrowHeadX + direction * arrowHeadSize * Math.cos(arrowHeadAngle + Math.PI / 6),
+                arrowHeadY + direction * arrowHeadSize * Math.sin(arrowHeadAngle + Math.PI / 6)
+            );
+            ctx.strokeStyle = torqueArrowColor;
+            ctx.lineWidth = 2/25*scaleX;
+            ctx.stroke();
+        }
+
+    }
+}
+        )RL_TOOLS_LITERAL";
+        return ui;
+    }
+    template <typename DEVICE, typename SPEC>
+    std::string get_description(DEVICE& device, rl::environments::multi_agent::Bottleneck<SPEC>& env){
+        std::string description;
+        description += "The agents have only local perception and need to squeeze through the bottleneck.";
+        return description;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif