450 lines
15 KiB
C++
450 lines
15 KiB
C++
//
|
||
// Created by lgv on 11/25/25.
|
||
//
|
||
|
||
|
||
#include <algorithm>
|
||
#include <cstdio>
|
||
#include <chrono>
|
||
#include <cmath>
|
||
#include "simulate/mujoco/mujoco_viewer/include/mujoco_viewer.h"
|
||
#include "simulate/mujoco/mujoco_viewer/include/array_safety.h"
|
||
#include "simulate/mujoco/mujoco_viewer/include/glfw_adapter.h"
|
||
|
||
namespace cmvr {
|
||
namespace mj = ::mujoco;
|
||
namespace mju = ::mujoco::sample_util;
|
||
|
||
using Seconds = std::chrono::duration<double>;
|
||
|
||
constexpr double kSyncMisalign = 0.1;
|
||
constexpr double kSimRefreshFraction = 0.7;
|
||
constexpr int kPiPMaxGeom = 100000;
|
||
|
||
class PiPGlfwAdapter : public mj::GlfwAdapter {
|
||
public:
|
||
explicit PiPGlfwAdapter(MuJocoViewer *owner)
|
||
: owner_(owner) {}
|
||
|
||
void SwapBuffers() override {
|
||
if (owner_) {
|
||
owner_->renderPiP();
|
||
}
|
||
mj::GlfwAdapter::SwapBuffers();
|
||
}
|
||
|
||
private:
|
||
MuJocoViewer *owner_;
|
||
};
|
||
|
||
|
||
static const char *CheckDiverged(int disableflags, const mjData *d) {
|
||
if (disableflags & mjDSBL_AUTORESET) {
|
||
for (mjtWarning w: {mjWARN_BADQACC, mjWARN_BADQVEL, mjWARN_BADQPOS}) {
|
||
if (d->warning[w].number > 0) {
|
||
return mju_warningText(w, d->warning[w].lastinfo);
|
||
}
|
||
}
|
||
}
|
||
return nullptr;
|
||
}
|
||
|
||
|
||
mjModel *MuJocoViewer::LoadModelSimple(const char *file,
|
||
char *load_error,
|
||
int error_sz) {
|
||
load_error[0] = '\0';
|
||
if (!file || !file[0]) {
|
||
std::snprintf(load_error, error_sz, "empty model filename");
|
||
return nullptr;
|
||
}
|
||
|
||
mjModel *mnew = nullptr;
|
||
const char *dot = std::strrchr(file, '.');
|
||
|
||
if (dot && std::strcmp(dot, ".mjb") == 0) {
|
||
mnew = mj_loadModel(file, nullptr);
|
||
if (!mnew) {
|
||
std::snprintf(load_error, error_sz,
|
||
"could not load binary model: %s", file);
|
||
}
|
||
} else {
|
||
mnew = mj_loadXML(file, nullptr, load_error, error_sz);
|
||
if (load_error[0]) {
|
||
int len = std::strlen(load_error);
|
||
if (len > 0 && load_error[len - 1] == '\n') {
|
||
load_error[len - 1] = '\0';
|
||
}
|
||
}
|
||
}
|
||
return mnew;
|
||
}
|
||
|
||
MuJocoViewer::MuJocoViewer(const char *model_path)
|
||
: model_path_(model_path ? model_path : "") {
|
||
std::printf("MuJoCo version %s\n", mj_versionString());
|
||
if (mjVERSION_HEADER != mj_version()) {
|
||
mju_error("Headers and library have different versions");
|
||
}
|
||
|
||
initSim();
|
||
}
|
||
|
||
void MuJocoViewer::initSim() {
|
||
mjv_defaultCamera(&cam_);
|
||
mjv_defaultOption(&opt_);
|
||
mjv_defaultPerturb(&pert_);
|
||
mjv_defaultCamera(&pip_cam_);
|
||
mjv_defaultScene(&pip_scene_);
|
||
|
||
auto platform_ui = std::make_unique<PiPGlfwAdapter>(this);
|
||
sim_ = std::make_unique<mj::Simulate>(
|
||
std::move(platform_ui),
|
||
&cam_, &opt_, &pert_,
|
||
/*is_passive=*/false
|
||
);
|
||
}
|
||
|
||
MuJocoViewer::~MuJocoViewer() {
|
||
if (sim_) {
|
||
sim_->exitrequest.store(1);
|
||
}
|
||
if (physics_started_.load() && physics_thread_.joinable()) {
|
||
physics_thread_.join();
|
||
}
|
||
|
||
if (pip_scene_inited_) {
|
||
mjv_freeScene(&pip_scene_);
|
||
pip_scene_inited_ = false;
|
||
}
|
||
|
||
if (d_) mj_deleteData(d_);
|
||
if (m_) mj_deleteModel(m_);
|
||
}
|
||
|
||
void MuJocoViewer::physicsThreadFunc() {
|
||
// 1) 加载模型
|
||
char error[1024];
|
||
m_ = LoadModelSimple(model_path_.c_str(), error, sizeof(error));
|
||
if (!m_) {
|
||
mju_error("Load model failed: %s", error);
|
||
}
|
||
|
||
d_ = mj_makeData(m_);
|
||
if (!d_) {
|
||
mj_deleteModel(m_);
|
||
m_ = nullptr;
|
||
mju_error("mj_makeData failed");
|
||
}
|
||
|
||
// 2) 用官方的方式：先告诉 Simulate 要加载，然后等 RenderLoop 处理
|
||
sim_->Load(m_, d_, model_path_.c_str());
|
||
|
||
{
|
||
// official main 里也是在 Load 之后先 forward 一下
|
||
const std::unique_lock<std::recursive_mutex> lock(sim_->mtx);
|
||
mj_forward(m_, d_);
|
||
}
|
||
|
||
physics_started_.store(true);
|
||
|
||
// 3) 进入物理循环
|
||
physicsLoop();
|
||
}
|
||
|
||
void MuJocoViewer::getCameraState(double &distance,
|
||
double &azimuth,
|
||
double &elevation) const {
|
||
distance = cam_.distance;
|
||
azimuth = cam_.azimuth;
|
||
elevation = cam_.elevation;
|
||
}
|
||
|
||
void MuJocoViewer::setupCamera(double distance,
|
||
double azimuth,
|
||
double elevation) {
|
||
// 相机初始化：根据模型设置一个合理的 free camera
|
||
if (m_) {
|
||
mjv_defaultFreeCamera(m_, &cam_);
|
||
// 以模型中心为观察点
|
||
cam_.lookat[0] = m_->stat.center[0];
|
||
cam_.lookat[1] = m_->stat.center[1];
|
||
cam_.lookat[2] = m_->stat.center[2];
|
||
} else {
|
||
mjv_defaultCamera(&cam_);
|
||
}
|
||
|
||
cam_.type = mjCAMERA_FREE;
|
||
cam_.distance = distance; // 距离
|
||
cam_.azimuth = azimuth; // 方位角（度）
|
||
cam_.elevation = elevation; // 俯仰角（度）
|
||
}
|
||
|
||
void MuJocoViewer::enablePiPCamera(const char *camera_name,
|
||
int left,
|
||
int bottom,
|
||
int width,
|
||
int height) {
|
||
pip_enabled_ = true;
|
||
pip_camera_name_ = camera_name ? camera_name : "";
|
||
pip_camera_id_ = -1;
|
||
pip_left_ = left;
|
||
pip_bottom_ = bottom;
|
||
pip_width_ = width > 0 ? width : 320;
|
||
pip_height_ = height > 0 ? height : 240;
|
||
pip_custom_pos_ = true;
|
||
}
|
||
|
||
void MuJocoViewer::disablePiPCamera() {
|
||
pip_enabled_ = false;
|
||
}
|
||
|
||
|
||
|
||
void MuJocoViewer::renderPiP() {
|
||
if (!pip_enabled_ || !sim_ || !m_ || !d_) return;
|
||
if (pip_camera_name_.empty()) return;
|
||
|
||
if (pip_camera_id_ < 0) {
|
||
pip_camera_id_ = mj_name2id(m_, mjOBJ_CAMERA, pip_camera_name_.c_str());
|
||
if (pip_camera_id_ < 0) {
|
||
return;
|
||
}
|
||
}
|
||
|
||
auto [fb_width, fb_height] = sim_->platform_ui->GetFramebufferSize();
|
||
if (fb_width <= 0 || fb_height <= 0) return;
|
||
|
||
int left = 0;
|
||
int bottom = 0;
|
||
int width = 0;
|
||
int height = 0;
|
||
|
||
if (pip_custom_pos_) {
|
||
left = std::max(0, std::min(pip_left_, fb_width - 1));
|
||
bottom = std::max(0, std::min(pip_bottom_, fb_height - 1));
|
||
width = std::min(pip_width_, fb_width - left);
|
||
height = std::min(pip_height_, fb_height - bottom);
|
||
} else {
|
||
width = std::min(pip_width_, fb_width - 2 * pip_margin_);
|
||
height = std::min(pip_height_, fb_height - 2 * pip_margin_);
|
||
left = fb_width - pip_margin_ - width;
|
||
bottom = pip_margin_;
|
||
}
|
||
|
||
if (width <= 0 || height <= 0) return;
|
||
|
||
mjrRect rect;
|
||
rect.width = width;
|
||
rect.height = height;
|
||
rect.left = left;
|
||
rect.bottom = bottom;
|
||
|
||
const std::unique_lock<std::recursive_mutex> lock(sim_->mtx);
|
||
|
||
if (!pip_scene_inited_ || pip_scene_model_ != m_) {
|
||
if (pip_scene_inited_) {
|
||
mjv_freeScene(&pip_scene_);
|
||
}
|
||
mjv_makeScene(m_, &pip_scene_, kPiPMaxGeom);
|
||
pip_scene_inited_ = true;
|
||
pip_scene_model_ = m_;
|
||
}
|
||
|
||
pip_cam_.type = mjCAMERA_FIXED;
|
||
pip_cam_.fixedcamid = pip_camera_id_;
|
||
pip_cam_.trackbodyid = -1;
|
||
|
||
mjv_updateScene(m_, d_, &opt_, &pert_, &pip_cam_, mjCAT_ALL, &pip_scene_);
|
||
mjr_render(rect, &pip_scene_, &sim_->platform_ui->mjr_context());
|
||
|
||
{
|
||
std::lock_guard<std::mutex> lock(pip_rgb_mtx_);
|
||
const int w = rect.width;
|
||
const int h = rect.height;
|
||
if (w > 0 && h > 0) {
|
||
pip_rgb_.resize(static_cast<size_t>(3 * w * h));
|
||
pip_depth_.resize(static_cast<size_t>(w * h));
|
||
|
||
// 同时读 RGB 和 depth(z-buffer 0..1)
|
||
mjr_readPixels(pip_rgb_.data(), pip_depth_.data(),
|
||
rect, &sim_->platform_ui->mjr_context());
|
||
|
||
// OpenGL 像素原点在左下，需要竖直翻转 RGB 和 depth
|
||
for (int r = 0; r < h / 2; ++r) {
|
||
// flip rgb row
|
||
unsigned char *top_row = pip_rgb_.data() + 3 * w * r;
|
||
unsigned char *bottom_row = pip_rgb_.data() + 3 * w * (h - 1 - r);
|
||
std::swap_ranges(top_row, top_row + 3 * w, bottom_row);
|
||
|
||
// flip depth row
|
||
float *top_d = pip_depth_.data() + w * r;
|
||
float *bot_d = pip_depth_.data() + w * (h - 1 - r);
|
||
std::swap_ranges(top_d, top_d + w, bot_d);
|
||
}
|
||
|
||
pip_rgb_width_ = w;
|
||
pip_rgb_height_ = h;
|
||
pip_rgb_valid_ = true;
|
||
++pip_frame_id_; // 新帧
|
||
}
|
||
}
|
||
}
|
||
|
||
void MuJocoViewer::physicsLoop() {
|
||
using Clock = mj::Simulate::Clock;
|
||
|
||
std::chrono::time_point<Clock> syncCPU;
|
||
mjtNum syncSim = 0;
|
||
|
||
while (!sim_->exitrequest.load()) {
|
||
// 给 UI 线程一点时间
|
||
if (sim_->run && sim_->busywait) {
|
||
std::this_thread::yield();
|
||
} else {
|
||
std::this_thread::sleep_for(std::chrono::milliseconds(1));
|
||
}
|
||
|
||
{
|
||
const std::unique_lock<std::recursive_mutex> lock(sim_->mtx);
|
||
if (!m_ || !d_) continue;
|
||
|
||
if (sim_->run) {
|
||
bool stepped = false;
|
||
const auto startCPU = Clock::now();
|
||
const auto elapsedCPU = startCPU - syncCPU;
|
||
double elapsedSim = d_->time - syncSim;
|
||
|
||
double slowdown = 100.0 / sim_->percentRealTime[sim_->real_time_index];
|
||
bool misaligned =
|
||
std::abs(Seconds(elapsedCPU).count() / slowdown - elapsedSim) > kSyncMisalign;
|
||
|
||
if (elapsedSim < 0 || elapsedCPU.count() < 0 ||
|
||
syncCPU.time_since_epoch().count() == 0 ||
|
||
misaligned || sim_->speed_changed) {
|
||
syncCPU = startCPU;
|
||
syncSim = d_->time;
|
||
sim_->speed_changed = false;
|
||
|
||
sim_->InjectNoise(sim_->key);
|
||
|
||
|
||
controlDispatch(m_, d_);
|
||
|
||
mj_step(m_, d_);
|
||
const char *msg = CheckDiverged(m_->opt.disableflags, d_);
|
||
if (msg) {
|
||
sim_->run = 0;
|
||
mju::strcpy_arr(sim_->load_error, msg);
|
||
} else {
|
||
stepped = true;
|
||
}
|
||
} else {
|
||
bool measured = false;
|
||
mjtNum prevSim = d_->time;
|
||
double refreshTime = kSimRefreshFraction / sim_->refresh_rate;
|
||
|
||
while (Seconds((d_->time - syncSim) * slowdown) <
|
||
(Clock::now() - syncCPU) &&
|
||
(Clock::now() - startCPU) < Seconds(refreshTime)) {
|
||
if (!measured && elapsedSim) {
|
||
sim_->measured_slowdown =
|
||
Seconds(elapsedCPU).count() / elapsedSim;
|
||
measured = true;
|
||
}
|
||
|
||
sim_->InjectNoise(sim_->key);
|
||
|
||
|
||
controlDispatch(m_, d_);
|
||
|
||
mj_step(m_, d_);
|
||
const char *msg = CheckDiverged(m_->opt.disableflags, d_);
|
||
if (msg) {
|
||
sim_->run = 0;
|
||
mju::strcpy_arr(sim_->load_error, msg);
|
||
} else {
|
||
stepped = true;
|
||
}
|
||
|
||
if (d_->time < prevSim) {
|
||
break;
|
||
}
|
||
}
|
||
}
|
||
|
||
if (stepped) {
|
||
sim_->AddToHistory();
|
||
}
|
||
} else {
|
||
mj_forward(m_, d_);
|
||
if (sim_->pause_update) {
|
||
mju_copy(d_->qacc_warmstart, d_->qacc, m_->nv);
|
||
}
|
||
sim_->speed_changed = true;
|
||
}
|
||
|
||
if (d_->time == 0.0) {
|
||
onReset(m_, d_);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
void MuJocoViewer::controlDispatch(mjModel *m, mjData *d) {
|
||
// 第一次：只跑一次 initOnce
|
||
if (!inited_) {
|
||
initOnce(m, d);
|
||
inited_ = true;
|
||
last_time_ = d->time;
|
||
}
|
||
|
||
// 检测 reset：时间从大跳到小（比如从 3.2 回到 0）
|
||
if (d->time < last_time_) {
|
||
onReset(m, d);
|
||
inited_ = false;
|
||
}
|
||
last_time_ = d->time;
|
||
|
||
// 每步控制
|
||
controlCallback(m, d);
|
||
}
|
||
|
||
void MuJocoViewer::run() {
|
||
// 先起 physics 线程（里面会 sim_->Load）
|
||
physics_thread_ = std::thread(&MuJocoViewer::physicsThreadFunc, this);
|
||
|
||
// 当前线程跑 UI + 渲染（阻塞）
|
||
sim_->RenderLoop();
|
||
|
||
// UI 退出后，通知 physics 线程结束
|
||
sim_->exitrequest.store(2);
|
||
|
||
if (physics_thread_.joinable()) {
|
||
physics_thread_.join();
|
||
}
|
||
}
|
||
|
||
|
||
uint64_t MuJocoViewer::getPiPCameraFrameId() const {
|
||
std::lock_guard<std::mutex> lock(pip_rgb_mtx_);
|
||
return pip_frame_id_;
|
||
}
|
||
|
||
bool MuJocoViewer::getPiPCameraRGBD(std::vector<unsigned char> &rgb,
|
||
std::vector<float> &depth,
|
||
int &width,
|
||
int &height,
|
||
uint64_t &frame_id) const {
|
||
std::lock_guard<std::mutex> lock(pip_rgb_mtx_);
|
||
if (!pip_rgb_valid_ || pip_rgb_.empty()) return false;
|
||
|
||
rgb = pip_rgb_;
|
||
depth = pip_depth_;
|
||
width = pip_rgb_width_;
|
||
height = pip_rgb_height_;
|
||
frame_id = pip_frame_id_;
|
||
return true;
|
||
}
|
||
} // namespace cmvr
|