Пример Mujoco + GLFW со слайдерами пространства суставов с использованием Dear imgui

git submodule add https://github.com/ocornut/imgui.git

#include <mujoco/mujoco.h>
#include <GLFW/glfw3.h>
#include <GL/gl.h>
#include "imgui/imgui.h"
#include "imgui/backends/imgui_impl_glfw.h"
#include "imgui/backends/imgui_impl_opengl2.h"
#include <cstdio>
#include <cstdlib>
#include <cmath>
#include <time.h>
#include <vector>
#include <string>
#include <algorithm>

// --- Глобальные переменные, минимальные для callback-функций GLFW ---
static mjModel* m = nullptr;
static mjData*  d = nullptr;

static mjvCamera cam;       // абстрактная камера
static mjvOption opt;       // опции визуализации
static mjvScene  scn;       // абстрактная сцена
static mjrContext con;      // пользовательский GPU-контекст

static bool  button_left = false, button_middle = false, button_right = false;
static double lastx = 0, lasty = 0;
// Синхронизация реального времени: запись времён начала симуляции и wall-clock
static double sim_start_time = 0.0;
static double realtime_start = 0.0;

// --- Простое состояние UI для слайдеров ---
struct SliderDef {
  int jnt;           // индекс сустава, к которому относится эта группа слайдеров
  int comp;          // компонент в суставе: -1 для hinge/slide, 0..2 Euler для ball, 0..5 (xyz,rpy) для free
  double vmin;       // минимальное значение слайдера (единицы: м или рад)
  double vmax;       // максимальное значение слайдера
  double value;      // текущее значение
  std::string label; // метка (не отображается как текст в этом минимальном UI)
};

static std::vector<SliderDef> g_sliders;   // все слайдеры

// Утилиты: clamp
template <typename T>
static inline T clamp(T v, T lo, T hi) { return std::max(lo, std::min(v, hi)); }

// Помощники для кватернионов (конвенция ZYX yaw-pitch-roll)
static void quat_normalize(double q[4]) {
  double n = std::sqrt(q[0]*q[0] + q[1]*q[1] + q[2]*q[2] + q[3]*q[3]);
  if (n <= 0) { q[0]=1; q[1]=q[2]=q[3]=0; return; }
  q[0]/=n; q[1]/=n; q[2]/=n; q[3]/=n;
}

static void euler_rpy_to_quat(double roll, double pitch, double yaw, double q[4]) {
  double cr = std::cos(roll*0.5),  sr = std::sin(roll*0.5);
  double cp = std::cos(pitch*0.5), sp = std::sin(pitch*0.5);
  double cy = std::cos(yaw*0.5),   sy = std::sin(yaw*0.5);
  // Порядок ZYX: q = Rz(yaw)*Ry(pitch)*Rx(roll)
  q[0] = cr*cp*cy + sr*sp*sy; // w
  q[1] = sr*cp*cy - cr*sp*sy; // x
  q[2] = cr*sp*cy + sr*cp*sy; // y
  q[3] = cr*cp*sy - sr*sp*cy; // z
  quat_normalize(q);
}

static void quat_to_euler_rpy(const double qin[4], double& roll, double& pitch, double& yaw) {
  double q[4] = {qin[0], qin[1], qin[2], qin[3]};
  quat_normalize(q);
  // roll (вращение вокруг оси x)
  double sinr_cosp = 2.0*(q[0]*q[1] + q[2]*q[3]);
  double cosr_cosp = 1.0 - 2.0*(q[1]*q[1] + q[2]*q[2]);
  roll = std::atan2(sinr_cosp, cosr_cosp);
  // pitch (ось y)
  double sinp = 2.0*(q[0]*q[2] - q[3]*q[1]);
  if (std::abs(sinp) >= 1.0) pitch = std::copysign(M_PI/2.0, sinp);
  else pitch = std::asin(sinp);
  // yaw (ось z)
  double siny_cosp = 2.0*(q[0]*q[3] + q[1]*q[2]);
  double cosy_cosp = 1.0 - 2.0*(q[2]*q[2] + q[3]*q[3]);
  yaw = std::atan2(siny_cosp, cosy_cosp);
}

// Построение слайдеров из модели и текущих данных
static void build_sliders() {
  g_sliders.clear();
  if (!m || !d) return;
  for (int j=0; j<m->njnt; ++j) {
    int type = m->jnt_type[j];
    int qadr = m->jnt_qposadr[j];
    const char* name = (m->names && m->name_jntadr) ? m->names + m->name_jntadr[j] : "joint";
    if (type == mjJNT_HINGE) {
      double vmin = -M_PI, vmax = M_PI;
      if (m->jnt_limited[j]) { vmin = m->jnt_range[2*j]; vmax = m->jnt_range[2*j+1]; }
      g_sliders.push_back({j, -1, vmin, vmax, d->qpos[qadr], std::string(name)});
    } else if (type == mjJNT_SLIDE) {
      double vmin = -1.0, vmax = 1.0;
      if (m->jnt_limited[j]) { vmin = m->jnt_range[2*j]; vmax = m->jnt_range[2*j+1]; }
      g_sliders.push_back({j, -1, vmin, vmax, d->qpos[qadr], std::string(name)});
    } else if (type == mjJNT_BALL) {
      double q[4] = { d->qpos[qadr+0], d->qpos[qadr+1], d->qpos[qadr+2], d->qpos[qadr+3] };
      double r,p,y; quat_to_euler_rpy(q, r,p,y);
      g_sliders.push_back({j, 0, -M_PI, M_PI, r, std::string(name)+".roll"});
      g_sliders.push_back({j, 1, -M_PI, M_PI, p, std::string(name)+".pitch"});
      g_sliders.push_back({j, 2, -M_PI, M_PI, y, std::string(name)+".yaw"});
    } else if (type == mjJNT_FREE) {
      // позиция xyz
      g_sliders.push_back({j, 0, -2.0, 2.0, d->qpos[qadr+0], std::string(name)+".x"});
      g_sliders.push_back({j, 1, -2.0, 2.0, d->qpos[qadr+1], std::string(name)+".y"});
      g_sliders.push_back({j, 2, -2.0, 2.0, d->qpos[qadr+2], std::string(name)+".z"});
      double q[4] = { d->qpos[qadr+3], d->qpos[qadr+4], d->qpos[qadr+5], d->qpos[qadr+6] };
      double r,p,y; quat_to_euler_rpy(q, r,p,y);
      g_sliders.push_back({j, 3, -M_PI, M_PI, r, std::string(name)+".roll"});
      g_sliders.push_back({j, 4, -M_PI, M_PI, p, std::string(name)+".pitch"});
      g_sliders.push_back({j, 5, -M_PI, M_PI, y, std::string(name)+".yaw"});
    }
  }
}

// Применение значений слайдеров к qpos и прямой кинематики
static void apply_sliders_to_qpos() {
  if (!m || !d) return;
  // Начинаем с текущего qpos
  // Для каждого сустава записываем значения
  for (size_t i=0; i<g_sliders.size(); ++i) {
    const SliderDef& s = g_sliders[i];
    int type = m->jnt_type[s.jnt];
    int qadr = m->jnt_qposadr[s.jnt];
    if (type == mjJNT_HINGE || type == mjJNT_SLIDE) {
      d->qpos[qadr] = s.value;
    } else if (type == mjJNT_BALL) {
      // Найти три слайдера для этого сустава
      double r=0,p=0,y=0;
      for (const auto& t : g_sliders) if (t.jnt==s.jnt) {
        if (t.comp==0) r=t.value; else if (t.comp==1) p=t.value; else if (t.comp==2) y=t.value;
      }
      double q[4]; euler_rpy_to_quat(r,p,y,q);
      d->qpos[qadr+0]=q[0]; d->qpos[qadr+1]=q[1]; d->qpos[qadr+2]=q[2]; d->qpos[qadr+3]=q[3];
    } else if (type == mjJNT_FREE) {
      double pos[3]; double r=0,p=0,y=0;
      for (const auto& t : g_sliders) if (t.jnt==s.jnt) {
        if (t.comp==0) pos[0]=t.value; else if (t.comp==1) pos[1]=t.value; else if (t.comp==2) pos[2]=t.value;
        else if (t.comp==3) r=t.value; else if (t.comp==4) p=t.value; else if (t.comp==5) y=t.value;
      }
      double q[4]; euler_rpy_to_quat(r,p,y,q);
      d->qpos[qadr+0]=pos[0]; d->qpos[qadr+1]=pos[1]; d->qpos[qadr+2]=pos[2];
      d->qpos[qadr+3]=q[0];  d->qpos[qadr+4]=q[1];  d->qpos[qadr+5]=q[2];  d->qpos[qadr+6]=q[3];
    }
  }
  mj_forward(m, d);
}

// Нет пользовательского оверлея; ImGui будет обрабатывать весь рендеринг UI и взаимодействия

// --- Callback ошибок GLFW ---
static void glfw_error_callback(int error, const char* description) {
  std::fprintf(stderr, "GLFW error %d: %s\n", error, description);
}

// --- Кнопка мыши: запоминаем, какие кнопки нажаты ---
static void mouse_button_callback(GLFWwindow* window, int button, int act, int /*mods*/) {
  // Простая, стабильная обработка состояния кнопок
  if (button == GLFW_MOUSE_BUTTON_LEFT)   button_left   = (act == GLFW_PRESS) ? true : (act == GLFW_RELEASE ? false : button_left);
  if (button == GLFW_MOUSE_BUTTON_MIDDLE) button_middle = (act == GLFW_PRESS) ? true : (act == GLFW_RELEASE ? false : button_middle);
  if (button == GLFW_MOUSE_BUTTON_RIGHT)  button_right  = (act == GLFW_PRESS) ? true : (act == GLFW_RELEASE ? false : button_right);

  // Запись текущей позиции курсора (оконные координаты)
  glfwGetCursorPos(window, &lastx, &lasty);
}

// --- Движение мыши: преобразование движения в действие камеры ---
static void cursor_pos_callback(GLFWwindow* window, double xpos, double ypos) {
  if (!m) return;
  // Если ни одна кнопка мыши не нажата, просто обновляем последнюю позицию
  if (!button_left && !button_right && !button_middle) {
    lastx = xpos;
    lasty = ypos;
    return;
  }

  // Если ImGui хочет мышь, не двигаем камеру
  if (ImGui::GetIO().WantCaptureMouse) {
    lastx = xpos; lasty = ypos;
    return;
  }

  // Вычисление движения в оконных координатах, нормализация по высоте framebuffer (стабильно при разных DPI)
  int fbw = 0, fbh = 1;
  glfwGetFramebufferSize(window, &fbw, &fbh);
  if (fbh <= 0) fbh = 1;

  double dx = (xpos - lastx) / (double)fbh;
  double dy = (ypos - lasty) / (double)fbh;
  lastx = xpos;
  lasty = ypos;

  if (button_left) {
    // Левый клик с перетаскиванием: вращение
    mjv_moveCamera(m, mjMOUSE_ROTATE_H, dx, dy, &scn, &cam);
  } else if (button_right) {
    // Правый клик с перетаскиванием: горизонтальное панорамирование
    mjv_moveCamera(m, mjMOUSE_MOVE_H, dx, dy, &scn, &cam);
  } else if (button_middle) {
    // Средний клик с перетаскиванием: масштабирование
    mjv_moveCamera(m, mjMOUSE_ZOOM, dx, dy, &scn, &cam);
  }
}

// --- Колесо прокрутки для масштабирования ---
static void scroll_callback(GLFWwindow* window, double /*xoffset*/, double yoffset) {
  if (!m) return;
  // Прокрутка для масштабирования (мягкий коэффициент)
  mjv_moveCamera(m, mjMOUSE_ZOOM, 0, -0.05 * yoffset, &scn, &cam);
}

// --- Клавиша: 'R' для сброса ---
static void key_callback(GLFWwindow* window, int key, int scancode, int act, int mods) {
  if (act == GLFW_PRESS || act == GLFW_REPEAT) {
    if (key == GLFW_KEY_R) {
      // Сброс состояния симуляции и таймеров реального времени, чтобы симуляция не убежала вперёд
      mj_resetData(m, d);
      sim_start_time = (d ? d->time : 0.0);
      realtime_start = glfwGetTime();
  // Перестроение и синхронизация слайдеров со сброшенным состоянием
  build_sliders();
  apply_sliders_to_qpos();
    }
    if (key == GLFW_KEY_ESCAPE) glfwSetWindowShouldClose(window, GLFW_TRUE);
  }
}

int main(int argc, char** argv) {
  // Использование
  if (argc < 2) {
    std::printf("Usage: %s model.xml\n", argv[0]);
    return 1;
  }

  // (Open-source MuJoCo >= 2.2 не требует активации лицензии)

  // Загрузка модели
  char error[1024] = {0};
  m = mj_loadXML(argv[1], nullptr, error, sizeof(error));
  if (!m) {
    std::fprintf(stderr, "Could not load model: %s\n", error);
    return 1;
  }
  d = mj_makeData(m);

  // Инициализация GLFW
  glfwSetErrorCallback(glfw_error_callback);
  if (!glfwInit()) {
    std::fprintf(stderr, "Could not initialize GLFW\n");
    return 1;
  }

  // Создание окна/контекста
  glfwWindowHint(GLFW_DOUBLEBUFFER, 1);
  GLFWwindow* window = glfwCreateWindow(1200, 900, "MuJoCo Minimal Viewer", nullptr, nullptr);
  if (!window) {
    std::fprintf(stderr, "Could not create GLFW window\n");
    glfwTerminate();
    return 1;
  }
  glfwMakeContextCurrent(window);
  glfwSwapInterval(0); // отключить vsync

  // Установка callback-функций
  glfwSetMouseButtonCallback(window, mouse_button_callback);
  glfwSetCursorPosCallback(window, cursor_pos_callback);
  glfwSetScrollCallback(window, scroll_callback);
  glfwSetKeyCallback(window, key_callback);

  // Инициализация визуализации MuJoCo
  mjv_defaultCamera(&cam);
  // Установка точки обзора камеры так, чтобы начальный вид был центрирован на модели
  cam.lookat[0] = 0.0;
  cam.lookat[1] = 0.0;
  cam.lookat[2] = 0.5;
  mjv_defaultOption(&opt);
  mjv_defaultScene(&scn);
  mjr_defaultContext(&con);

  mjv_makeScene(m, &scn, 2000);               // выделение сцены
  mjr_makeContext(m, &con, mjFONTSCALE_100);  // выделение GPU-контекста

  // Инициализация таймеров синхронизации реального времени (запуск обоих в одно и то же wall-clock время)
  sim_start_time = d->time;
  realtime_start = glfwGetTime();

  // Построение начальных слайдеров из модели
  build_sliders();
  apply_sliders_to_qpos();

  // --- Настройка ImGui ---
  IMGUI_CHECKVERSION();
  ImGui::CreateContext();
  ImGuiIO& io = ImGui::GetIO(); (void)io;
  ImGui::StyleColorsDark();
  ImGui_ImplGlfw_InitForOpenGL(window, true);
  ImGui_ImplOpenGL2_Init();

  // Размещение камеры для показа всей модели
  cam.type = mjCAMERA_FREE;
  mj_forward(m, d);

  // Главный цикл
  while (!glfwWindowShouldClose(window)) {
  // Симуляция отключена: напрямую используем qpos из слайдеров и поддерживаем mj_forward в актуальном состоянии
  // В случае, если что-то внешнее изменило d->qpos, мы можем повторно применять каждый кадр; достаточно дёшево
  apply_sliders_to_qpos();

    // Получение размера framebuffer и установка viewport
    int width, height;
    glfwGetFramebufferSize(window, &width, &height);
    mjrRect viewport = {0, 0, width, height};

    // Обновление сцены и рендеринг
    mjv_updateScene(m, d, &opt, nullptr, &cam, mjCAT_ALL, &scn);
    mjr_render(viewport, &scn, &con);

    // UI ImGui
    ImGui_ImplOpenGL2_NewFrame();
    ImGui_ImplGlfw_NewFrame();
    ImGui::NewFrame();

    ImGui::SetNextWindowPos(ImVec2(10, 10), 0);
    ImGui::SetNextWindowSize(ImVec2(350, std::min(700, height-20)), 0);
    if (ImGui::Begin("Joints", nullptr, 0)) {
      for (size_t i=0; i<g_sliders.size(); ++i) {
        SliderDef& s = g_sliders[i];
        float v = (float)s.value;
        float vmin = (float)s.vmin;
        float vmax = (float)s.vmax;
        std::string lab = s.label;
        if (ImGui::SliderFloat(lab.c_str(), &v, vmin, vmax)) {
          s.value = v;
        }
      }
    }
    ImGui::End();

    // Применение изменений слайдеров и рендеринг UI поверх
    apply_sliders_to_qpos();
    ImGui::Render();
    ImGui_ImplOpenGL2_RenderDrawData(ImGui::GetDrawData());

    // Swap и poll
    glfwSwapBuffers(window);
    glfwPollEvents();
  }

  // Очистка
  mjr_freeContext(&con);
  mjv_freeScene(&scn);
  mj_deleteData(d);
  mj_deleteModel(m);
  // Завершение ImGui
  ImGui_ImplOpenGL2_Shutdown();
  ImGui_ImplGlfw_Shutdown();
  ImGui::DestroyContext();
  glfwTerminate();
  return 0;
}

cmake_minimum_required(VERSION 3.10)
project(mujocotest)

set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
set(CMAKE_CXX_EXTENSIONS OFF)

find_package(PkgConfig REQUIRED)
pkg_search_module(MUJOCO mujoco)
pkg_search_module(GLFW REQUIRED glfw3)

# Предпочтение GLVND на системах, которые его поддерживают. Если вам нужен legacy-загрузчик, установите
# OpenGL_GL_PREFERENCE в "LEGACY" в командной строке cmake или в кэше.
set(OpenGL_GL_PREFERENCE "GLVND" CACHE STRING "Prefer GLVND or LEGACY for OpenGL loader")

find_package(OpenGL REQUIRED)

if(NOT MUJOCO_FOUND)
	set(MUJOCO_INCLUDE_DIRS /opt/mujoco/include)
	set(MUJOCO_LIBRARIES /opt/mujoco/lib/libmujoco.so)
endif()

add_executable(mujocotest main.cpp)
target_include_directories(mujocotest PRIVATE ${MUJOCO_INCLUDE_DIRS} ${GLFW_INCLUDE_DIRS} ${CMAKE_CURRENT_SOURCE_DIR}/imgui ${CMAKE_CURRENT_SOURCE_DIR}/imgui/backends)

set(IMGUI_SOURCES
	${CMAKE_CURRENT_SOURCE_DIR}/imgui/imgui.cpp
	${CMAKE_CURRENT_SOURCE_DIR}/imgui/imgui_draw.cpp
	${CMAKE_CURRENT_SOURCE_DIR}/imgui/imgui_tables.cpp
	${CMAKE_CURRENT_SOURCE_DIR}/imgui/imgui_widgets.cpp
	${CMAKE_CURRENT_SOURCE_DIR}/imgui/backends/imgui_impl_glfw.cpp
	${CMAKE_CURRENT_SOURCE_DIR}/imgui/backends/imgui_impl_opengl2.cpp
)

target_sources(mujocotest PRIVATE ${IMGUI_SOURCES})
target_link_libraries(mujocotest PRIVATE ${MUJOCO_LIBRARIES} ${GLFW_LIBRARIES} OpenGL::GL)

mkdir build
cmake .
make
./mujocotest ~/mujoco_menagerie/franka_fr3/fr3.xml