From 6f5a7576f9f7a528decfe3424a4a1f77b8de2cfb Mon Sep 17 00:00:00 2001 From: Haroon Qureshi Date: Thu, 18 Sep 2025 02:14:12 -0700 Subject: [PATCH] Move `simlib` into `experimental` folder and rename it `toolbox`. Also: - update the namespace to `toolbox` (instead of `filament`) - Add copyright notices to all files PiperOrigin-RevId: 808474353 Change-Id: Id1fe86c5d183c8e91bec8509b6c0830b5fe84c1c --- src/experimental/toolbox/helpers.cc | 231 ++++++++++++++++++ src/experimental/toolbox/helpers.h | 66 +++++ src/experimental/toolbox/imgui_widgets.cc | 92 +++++++ src/experimental/toolbox/imgui_widgets.h | 239 ++++++++++++++++++ src/experimental/toolbox/interaction.cc | 281 ++++++++++++++++++++++ src/experimental/toolbox/interaction.h | 62 +++++ src/experimental/toolbox/physics.cc | 222 +++++++++++++++++ src/experimental/toolbox/physics.h | 126 ++++++++++ src/experimental/toolbox/renderer.cc | 82 +++++++ src/experimental/toolbox/renderer.h | 82 +++++++ src/experimental/toolbox/step_control.cc | 178 ++++++++++++++ src/experimental/toolbox/step_control.h | 93 +++++++ src/experimental/toolbox/window.cc | 168 +++++++++++++ src/experimental/toolbox/window.h | 87 +++++++ 14 files changed, 2009 insertions(+) create mode 100644 src/experimental/toolbox/helpers.cc create mode 100644 src/experimental/toolbox/helpers.h create mode 100644 src/experimental/toolbox/imgui_widgets.cc create mode 100644 src/experimental/toolbox/imgui_widgets.h create mode 100644 src/experimental/toolbox/interaction.cc create mode 100644 src/experimental/toolbox/interaction.h create mode 100644 src/experimental/toolbox/physics.cc create mode 100644 src/experimental/toolbox/physics.h create mode 100644 src/experimental/toolbox/renderer.cc create mode 100644 src/experimental/toolbox/renderer.h create mode 100644 src/experimental/toolbox/step_control.cc create mode 100644 src/experimental/toolbox/step_control.h create mode 100644 src/experimental/toolbox/window.cc create mode 100644 src/experimental/toolbox/window.h diff --git a/src/experimental/toolbox/helpers.cc b/src/experimental/toolbox/helpers.cc new file mode 100644 index 00000000..29700f82 --- /dev/null +++ b/src/experimental/toolbox/helpers.cc @@ -0,0 +1,231 @@ +// Copyright 2025 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "experimental/toolbox/helpers.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "third_party/libwebp/src/webp/encode.h" +#include "third_party/libwebp/src/webp/types.h" +#include +#include +#include +#include "xml/xml_api.h" + +namespace mujoco::toolbox { + +mjModel* LoadMujocoModel(const std::string& model_file, const mjVFS* vfs) { + mjModel* model = nullptr; + + if (model_file.empty()) { + auto spec = mj_makeSpec(); + model = mj_compile(spec, 0); + mj_deleteSpec(spec); + } else if (model_file.ends_with(".mjb")) { + model = mj_loadModel(model_file.c_str(), 0); + if (!model) { + mju_error("Could not load binary model"); + } + } else if (model_file.ends_with(".xml")) { + char error[1000] = ""; + model = mj_loadXML(model_file.c_str(), vfs, error, sizeof(error)); + if (!model) { + mju_error("Load model error: %s", error); + } + } else { + char error[1000] = ""; + auto spec = + mj_parseXMLString(model_file.c_str(), nullptr, error, sizeof(error)); + if (!spec) { + mju_error("Load model error: %s", error); + } + model = mj_compile(spec, 0); + mj_deleteSpec(spec); + } + return model; +} + +void SaveText(const std::string& contents, const std::string& filename) { + std::ofstream file(filename); + file.write(contents.data(), contents.size()); + file.close(); +} + +std::string LoadText(const std::string& filename) { + std::ifstream file(filename); + std::string contents((std::istreambuf_iterator(file)), + std::istreambuf_iterator()); + file.close(); + return contents; +} + +void SaveColorToWebp(int width, int height, const unsigned char* data, + const std::string& filename) { + uint8_t* webp = nullptr; + const size_t size = + WebPEncodeLosslessRGB(data, width, height, width * 3, &webp); + + std::ofstream file(filename, std::ios::binary); + file.write(reinterpret_cast(webp), size); + file.close(); + WebPFree(webp); +} + +void SaveDepthToWebp(int width, int height, const float* data, + const std::string& filename) { + const int size = width * height; + + // Turn the depth buffer into a greyscale color buffer. + std::vector byte_buffer; + byte_buffer.reserve(size * 3); + for (int i = 0; i < size; ++i) { + auto byte = static_cast(255.0 * data[i]); + byte_buffer.push_back(byte); + byte_buffer.push_back(byte); + byte_buffer.push_back(byte); + } + SaveColorToWebp(width, height, byte_buffer.data(), filename); +} + +void SaveScreenshotToWebp(int width, int height, mjrContext* con, + const std::string& filename) { + mjr_setBuffer(mjFB_OFFSCREEN, con); + auto rgb_buffer = std::vector(3 * width * height); + auto depth_buffer = std::vector(width * height, 1.0f); + mjrRect viewport = {0, 0, width, height}; + mjr_readPixels(rgb_buffer.data(), depth_buffer.data(), viewport, con); + mjr_setBuffer(mjFB_WINDOW, con); + SaveColorToWebp(width, height, rgb_buffer.data(), filename); +} + +const void* GetValue(const mjModel* model, const mjData* data, + const char* field, int index) { + MJDATA_POINTERS_PREAMBLE(model); +#define X(TYPE, NAME, NR, NC) \ + if (!std::strcmp(#NAME, field) && !std::strcmp(#TYPE, "mjtNum")) { \ + if (index >= 0 && index < model->NR * NC) { \ + return &data->NAME[index]; \ + } else { \ + return nullptr; \ + } \ + } + MJDATA_POINTERS +#undef X + return nullptr; // Invalid field. +} + +std::string CameraToString(const mjvScene* scene) { + const mjvGLCamera* cameras = scene->camera; + const float pos_x = (cameras[0].pos[0] + cameras[1].pos[0]) / 2; + const float pos_y = (cameras[0].pos[1] + cameras[1].pos[1]) / 2; + const float pos_z = (cameras[0].pos[2] + cameras[1].pos[2]) / 2; + + mjtNum cam_forward[3]; + mju_f2n(cam_forward, cameras[0].forward, 3); + mjtNum cam_up[3]; + mju_f2n(cam_up, cameras[0].up, 3); + mjtNum cam_right[3]; + mju_cross(cam_right, cam_forward, cam_up); + + char str[500]; + std::snprintf(str, sizeof(str), + "\n", + pos_x, pos_y, pos_z, cam_right[0], cam_right[1], cam_right[2], + cam_up[0], cam_up[1], cam_up[2]); + return str; +} + +std::string KeyframeToString(const mjModel* model, const mjData* data, + bool full_precision) { + const int kStrLen = 5000; + + char buf[200]; + const char p_regular[] = "%g"; + const char p_full[] = "%-22.16g"; + const char* format = full_precision ? p_full : p_regular; + + char str[kStrLen] = "time); + std::strncat(str, buf, kStrLen); + + // qpos + std::strncat(str, "\"\n qpos=\"", kStrLen); + for (int i = 0; i < model->nq; i++) { + std::snprintf(buf, sizeof(buf), format, data->qpos[i]); + if (i < model->nq - 1) std::strncat(buf, " ", 200); + std::strncat(str, buf, kStrLen); + } + + // qvel + std::strncat(str, "\"\n qvel=\"", kStrLen); + for (int i = 0; i < model->nv; i++) { + std::snprintf(buf, sizeof(buf), format, data->qvel[i]); + if (i < model->nv - 1) std::strncat(buf, " ", 200); + std::strncat(str, buf, kStrLen); + } + + // act + if (model->na > 0) { + std::strncat(str, "\"\n act=\"", kStrLen); + for (int i = 0; i < model->na; i++) { + std::snprintf(buf, sizeof(buf), format, data->act[i]); + if (i < model->na - 1) std::strncat(buf, " ", 200); + std::strncat(str, buf, kStrLen); + } + } + + // ctrl + if (model->nu > 0) { + std::strncat(str, "\"\n ctrl=\"", kStrLen); + for (int i = 0; i < model->nu; i++) { + std::snprintf(buf, sizeof(buf), format, data->ctrl[i]); + if (i < model->nu - 1) std::strncat(buf, " ", 200); + std::strncat(str, buf, kStrLen); + } + } + + if (model->nmocap > 0) { + std::strncat(str, "\"\n mpos=\"", kStrLen); + for (int i = 0; i < 3 * model->nmocap; i++) { + std::snprintf(buf, sizeof(buf), format, data->mocap_pos[i]); + if (i < 3 * model->nmocap - 1) std::strncat(buf, " ", 200); + std::strncat(str, buf, kStrLen); + } + + // mocap_quat + std::strncat(str, "\"\n mquat=\"", kStrLen); + for (int i = 0; i < 4 * model->nmocap; i++) { + std::snprintf(buf, sizeof(buf), format, data->mocap_quat[i]); + if (i < 4 * model->nmocap - 1) std::strncat(buf, " ", 200); + std::strncat(str, buf, kStrLen); + } + } + + std::strncat(str, "\"\n/>", kStrLen); + return str; +} + +} // namespace mujoco::toolbox diff --git a/src/experimental/toolbox/helpers.h b/src/experimental/toolbox/helpers.h new file mode 100644 index 00000000..b2382c20 --- /dev/null +++ b/src/experimental/toolbox/helpers.h @@ -0,0 +1,66 @@ +// Copyright 2025 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Standalone functions used by Simulate. +#ifndef MUJOCO_SRC_EXPERIMENTAL_TOOLBOX_HELPERS_H_ +#define MUJOCO_SRC_EXPERIMENTAL_TOOLBOX_HELPERS_H_ + +#include +#include +#include +#include +#include + +#include +#include + +namespace mujoco::toolbox { + +// Function signature for loading assets from a given path. +using LoadAssetFn = std::function(std::string_view)>; + +// Save/load for simple ascii files. +void SaveText(const std::string& contents, const std::string& filename); +std::string LoadText(const std::string& filename); + +// Exports the given color buffer to a webp file. +void SaveColorToWebp(int width, int height, const unsigned char* data, + const std::string& filename); + +// Exports the given depth buffer to a webp file. +void SaveDepthToWebp(int width, int height, const float* data, + const std::string& filename); + +// Exports the current state of the mjrContext to a webp file. +void SaveScreenshotToWebp(int width, int height, mjrContext* con, + const std::string& filename); + +// Loads a MuJoCo model from the given file. +mjModel* LoadMujocoModel(const std::string& model_file, const mjVFS* vfs); + +// Returns a pointer to the value of the given field in the given data. +// Returns nullptr if the field is not found or the index is out of bounds. +const void* GetValue(const mjModel* model, const mjData* data, + const char* field, int index); + +// Returns an XML string representation of the scene cameras. +std::string CameraToString(const mjvScene* scene); + +// Returns an XML string representation of current data keyframe. +std::string KeyframeToString(const mjModel* model, const mjData* data, + bool full_precision = false); + +} // namespace mujoco::toolbox + +#endif // MUJOCO_SRC_EXPERIMENTAL_TOOLBOX_HELPERS_H_ diff --git a/src/experimental/toolbox/imgui_widgets.cc b/src/experimental/toolbox/imgui_widgets.cc new file mode 100644 index 00000000..40395b0a --- /dev/null +++ b/src/experimental/toolbox/imgui_widgets.cc @@ -0,0 +1,92 @@ +// Copyright 2025 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "experimental/toolbox/imgui_widgets.h" + +#include +#include +#include + +#include "third_party/dear_imgui/imgui.h" +#include + +namespace mujoco::toolbox { + +void AppendIniSection(std::string& ini, const std::string& section, + const KeyValues& key_values) { + if (section.front() != '[' || section.back() != ']') { + mju_error("Section must be enclosed in square brackets."); + } + ini += "\n" + std::string(section) + "\n"; + for (auto& [key, value] : key_values) { + ini += key + "=" + value + "\n"; + } +} + +KeyValues ReadIniSection(const std::string& contents, + const std::string& section) { + if (section.front() != '[' || section.back() != ']') { + mju_error("Section must be enclosed in square brackets."); + } + bool in_section = false; + + KeyValues key_values; + std::istringstream f(contents); + std::string line; + while (std::getline(f, line)) { + if (line[0] == '[') { + in_section = (line == section); + } else if (in_section) { + std::string::size_type pos = line.find('='); + if (pos != std::string::npos) { + key_values[line.substr(0, pos)] = line.substr(pos + 1); + } + } + } + return key_values; +} + +bool ImGui_Slider(const char* name, mjtNum* value, mjtNum min, mjtNum max) { + float f = *value; + const bool res = ImGui::SliderFloat(name, &f, min, max); + if (res) { + *value = f; + } + return res; +} + +bool ImGui_FileDialog(char* buf, int len) { + bool ok = false; + ImGui::Text("Filename"); + ImGui::SameLine(); + ImGui::InputText("##Filename", buf, len); + if (ImGui::Button("OK", ImVec2(120, 0))) { + ok = true; + ImGui::CloseCurrentPopup(); + } + ImGui::SetItemDefaultFocus(); + ImGui::SameLine(); + if (ImGui::Button("Cancel", ImVec2(120, 0))) { + ImGui::CloseCurrentPopup(); + } + return ok; +} + +void MaybeSaveToClipboard(const std::string& contents) { + if (ImGui::GetIO().SetClipboardTextFn) { + ImGui::GetIO().SetClipboardTextFn(nullptr, contents.c_str()); + } +} + +} // namespace mujoco::toolbox diff --git a/src/experimental/toolbox/imgui_widgets.h b/src/experimental/toolbox/imgui_widgets.h new file mode 100644 index 00000000..3595cc9f --- /dev/null +++ b/src/experimental/toolbox/imgui_widgets.h @@ -0,0 +1,239 @@ +// Copyright 2025 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MUJOCO_SRC_EXPERIMENTAL_TOOLBOX_IMGUI_WIDGETS_H_ +#define MUJOCO_SRC_EXPERIMENTAL_TOOLBOX_IMGUI_WIDGETS_H_ + +#include +#include +#include +#include +#include + +#include "third_party/dear_imgui/imgui.h" +#include "third_party/dear_imgui/imgui_internal.h" // For ButtonEx and PressedOnClick +#include + +namespace mujoco::toolbox { + +using KeyValues = std::unordered_map; + +// Appends key/value pairs to an Ini file. +void AppendIniSection(std::string& ini, const std::string& section, + const KeyValues& key_values); + +// Reads key/value pairs from an Ini file section. +KeyValues ReadIniSection(const std::string& contents, + const std::string& section); + +// ImGui file dialog. +bool ImGui_FileDialog(char* buf, int len); + +// ImGui Slider that supports both float and double types. +bool ImGui_Slider(const char* name, mjtNum* value, mjtNum min, mjtNum max); + +template +bool ImGui_Checkbox(const char* name, T& value) { + static_assert(std::is_integral()); + bool b = (value != 0); + const bool res = ImGui::Checkbox(name, &b); + if (res) { + value = b ? 1 : 0; + } + return res; +} + +enum class ToggleKind { + // Solid when ON and transparent when OFF. + kButton, + + // Slider which is right when ON and left when OFF. + kSlider, +}; + +template +bool Toggle(const char* label, T& boolean, + ToggleKind kind = ToggleKind::kButton, bool set_width = true) { + static_assert(std::is_integral_v, "Toggle only supports integral types."); + + // Compute this width once and cache it. Only used when set_width is true. + static int toggle_width = []() { + int longest = 0; + const char* longest_label = ""; + for (int i = 0; i < mjNVISFLAG; ++i) { + int length = static_cast(strlen(mjVISSTRING[i][0])); + if (length > longest) { + longest_label = mjVISSTRING[i][0]; + longest = length; + } + } + return ImGui::CalcTextSize(longest_label).x + 5; + }(); + + ImGui::PushID(label); + bool changed = false; + switch (kind) { + case ToggleKind::kButton: { + bool b = (boolean != 0); + bool transparent = !b; + + // NOTE(matijak): Its nice to have the button trigger on click but this + // requires using the currently internal PressedOnClick flag and ButtonEx + // function. It looks like this API has been stable for a long time, but + // in case it changes in a way which breaks and is annoying to maintain we + // can revert to the else clause and remove the imgui_internal.h include. + // Note that the else clause overrides different style colors since the UI + // is more intuitive with different settings. + if constexpr (true) { + ImColor button = ImGui::GetStyle().Colors[ImGuiCol_Button]; + if (transparent) button.Value.w = 0.0f; + ImGui::PushStyleColor(ImGuiCol_Button, (ImU32)button); + ImGui::PushStyleColor(ImGuiCol_ButtonActive, (ImU32)button); + ImGui::PushStyleColor(ImGuiCol_ButtonHovered, (ImU32)button); + + // Button width is set via an explicit size parameter, not via the + // SetNextItemWidth function. + ImVec2 size = set_width ? ImVec2(toggle_width, 0) : ImVec2(0, 0); + changed = ImGui::ButtonEx(label, size, ImGuiButtonFlags_PressedOnClick); + + if (changed) { + b = !b; + } + boolean = b; + + ImGui::PopStyleColor(3); + } else { + if (transparent) { + ImColor button = ImGui::GetStyle().Colors[ImGuiCol_Button]; + button.Value.w = 0.0f; + ImGui::PushStyleColor(ImGuiCol_Button, (ImU32)button); + } + + // Button width is set via an explicit size parameter, not via the + // SetNextItemWidth function. + ImVec2 size = set_width ? ImVec2(toggle_width, 0) : ImVec2(0, 0); + changed = ImGui::Button(label, size); + + if (changed) { + b = !b; + } + boolean = b; + + if (transparent) { + ImGui::PopStyleColor(1); + } + } + } break; + + case ToggleKind::kSlider: { + int i = (int)boolean; + const char* labels[2] = {label, label}; + const ImGuiSliderFlags flags = ImGuiSliderFlags_NoInput; + if (set_width) ImGui::SetNextItemWidth(toggle_width); + changed = ImGui::SliderInt("", &i, 0, 1, labels[i], flags); + boolean = (i != 0); + } break; + } + ImGui::PopID(); + return changed; +} + +inline bool ToggleBit(const char* label, int& flags, int flags_value, + ToggleKind kind = ToggleKind::kButton, + bool set_width = true) { + bool boolean = flags & flags_value; + bool changed = Toggle(label, boolean, kind, set_width); + if (changed) { + flags = boolean ? (flags | flags_value) : (flags & ~flags_value); + } + return changed; +} + +// Options for ImGui_InputN (see below). +template +struct ImGuiOpts { + std::optional min; + std::optional max; + std::optional step; + std::optional step_fast; + std::optional width; + const char* format = std::is_floating_point_v ? "%.3g" : "%d"; +}; + +// A compile-time wrapper around ImGui::InputScalarN. This is useful because +// MuJoCo uses an `mjtNum` type which is an alias for float or double. +// +// Options can be used to specify step sizes, clamp ranges, and formatting. +template +bool ImGui_InputN(const char* name, T* value, int num, ImGuiOpts opts = {}) { + bool res = false; + if (opts.width) { + ImGui::SetNextItemWidth(opts.width.value()); + } + if constexpr (std::is_same_v) { + const int step = opts.step.value_or(1); + const int step_fast = opts.step_fast.value_or(100); + const char* format = opts.format; + res = ImGui::InputScalarN(name, ImGuiDataType_S32, value, num, &step, + &step_fast, format); + + } else if constexpr (std::is_same_v) { + const float step = opts.step.value_or(0.f); + const float step_fast = opts.step_fast.value_or(0.f); + const float* pstep = opts.step.has_value() ? &step : nullptr; + const float* pstep_fast = opts.step_fast.has_value() ? &step_fast : nullptr; + const char* format = opts.format ? opts.format : "%.3f"; + res = ImGui::InputScalarN(name, ImGuiDataType_Float, value, num, pstep, + pstep_fast, format); + + } else if constexpr (std::is_same_v) { + const double step = opts.step.value_or(0.0); + const double step_fast = opts.step_fast.value_or(0.0); + const double* pstep = opts.step.has_value() ? &step : nullptr; + const double* pstep_fast = + opts.step_fast.has_value() ? &step_fast : nullptr; + const char* format = opts.format ? opts.format : "%.3f"; + res = ImGui::InputScalarN(name, ImGuiDataType_Double, value, num, pstep, + pstep_fast, format); + } else { + static_assert(false, "Unsupported type"); + } + + if (opts.min.has_value()) { + if (*value < *opts.min) *value = *opts.min; + } + if (opts.max.has_value()) { + if (*value > *opts.max) *value = *opts.max; + } + return res; +} + +template +bool ImGui_Input(const char* name, T* value, ImGuiOpts opts = {}) { + return ImGui_InputN(name, value, 1, opts); +} + +// Returns true if the given chord is has _just_ been pressed in this frame. +// (This is opposed to "Pressed" which means the chord is active, i.e. the user +// is holding down the keys.) +inline bool ImGui_IsChordJustPressed(ImGuiKeyChord chord) { + return ImGui::IsKeyChordPressed(chord, 0); +} + +// Saves the given contents to the clipboard if the clipboard is available. +void MaybeSaveToClipboard(const std::string& contents); + +} // namespace mujoco::toolbox + +#endif // MUJOCO_SRC_EXPERIMENTAL_TOOLBOX_IMGUI_WIDGETS_H_ diff --git a/src/experimental/toolbox/interaction.cc b/src/experimental/toolbox/interaction.cc new file mode 100644 index 00000000..7c78bbbb --- /dev/null +++ b/src/experimental/toolbox/interaction.cc @@ -0,0 +1,281 @@ +// Copyright 2025 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "experimental/toolbox/interaction.h" +#include + +#include "third_party/dear_imgui/imgui.h" +#include "experimental/toolbox/imgui_widgets.h" +#include "experimental/toolbox/physics.h" +#include "experimental/toolbox/renderer.h" +#include "experimental/toolbox/window.h" +#include + +namespace mujoco::toolbox { + +static void ToggleFlag(mjtByte& flag) { flag = flag ? 0 : 1; } + +static void SelectParentPerturb(const mjModel* model, mjvPerturb& perturb) { + if (perturb.select > 0) { + perturb.select = model->body_parentid[perturb.select]; + perturb.flexselect = -1; + perturb.skinselect = -1; + if (perturb.select <= 0) { + perturb.active = 0; + } + } +} + +PickResult Pick(float x, float y, Window* window, Renderer* renderer, + Physics* physics, const mjvOption& vis_options) { + const float w = static_cast(window->GetWidth()); + const float h = static_cast(window->GetHeight()); + const float aspect_ratio = w / h; + + PickResult result; + result.body = + mjv_select(physics->GetModel(), physics->GetData(), &vis_options, + aspect_ratio, x, 1.0f - y, &renderer->GetScene(), result.point, + &result.geom, &result.flex, &result.skin); + return result; +} + +int SetCamera(const mjModel& model, mjvCamera& camera, int request_idx) { + // 0 = free, 1 = tracking, 2+ = fixed + int camera_idx = std::clamp(request_idx, 0, std::max(model.ncam + 1, 0)); + if (camera_idx == 0) { + camera.type = mjCAMERA_FREE; + } else if (camera_idx == 1) { + if (camera.trackbodyid >= 0) { + camera.type = mjCAMERA_TRACKING; + camera.fixedcamid = -1; + } else { + camera.type = mjCAMERA_FREE; + camera_idx = 0; + } + } else { + camera.type = mjCAMERA_FIXED; + camera.fixedcamid = camera_idx - 2; + } + + return camera_idx; +} + +void HandleMouseEvents(Window* window, Renderer* renderer, + Physics* physics, mjvPerturb& perturb, + mjvOption& vis_options, mjvCamera& camera, + int& camera_idx) { + auto& io = ImGui::GetIO(); + + if (io.WantCaptureMouse) { + return; + } + + // Normalize mouse positions and movement to display size. + const float mouse_x = io.MousePos.x / io.DisplaySize.x; + const float mouse_y = io.MousePos.y / io.DisplaySize.y; + const float mouse_dx = io.MouseDelta.x / io.DisplaySize.x; + const float mouse_dy = io.MouseDelta.y / io.DisplaySize.y; + const float mouse_scroll = io.MouseWheel / 50.0f; + + mjModel* model = physics->GetModel(); + mjData* data = physics->GetData(); + mjvScene& scene = renderer->GetScene(); + + // Determine the mouse action based on which buttons are down. + mjtMouse action = mjMOUSE_NONE; + if (ImGui::IsMouseDown(ImGuiMouseButton_Left)) { + action = io.KeyShift ? mjMOUSE_ROTATE_H : mjMOUSE_ROTATE_V; + } else if (ImGui::IsMouseDown(ImGuiMouseButton_Right)) { + action = io.KeyShift ? mjMOUSE_MOVE_H : mjMOUSE_MOVE_V; + } else if (ImGui::IsMouseDown(ImGuiMouseButton_Middle)) { + action = mjMOUSE_ZOOM; + } else { + // If no mouse buttons are down, end any active perturbations. + perturb.active = 0; + } + + // Mouse scroll. + if (model && mouse_scroll != 0.0f) { + mjv_moveCamera(model, mjMOUSE_ZOOM, 0, mouse_scroll, &scene, &camera); + } + + // Mouse drag. + if (model && data && action != mjMOUSE_NONE && + (mouse_dx != 0.0f || mouse_dy != 0.0f)) { + // If ctrl is pressed, move the perturbation, otherwise move the camera. + if (io.KeyCtrl) { + if (perturb.select > 0) { + const int active = + action == mjMOUSE_MOVE_V ? mjPERT_TRANSLATE : mjPERT_ROTATE; + if (active != perturb.active) { + mjv_initPerturb(model, data, &scene, &perturb); + perturb.active = active; + } + mjv_movePerturb(model, data, action, mouse_dx, mouse_dy, &scene, + &perturb); + } + } else { + mjv_moveCamera(model, action, mouse_dx, mouse_dy, &scene, &camera); + } + } + + // Left double click. + if (data && ImGui::IsMouseDoubleClicked(ImGuiMouseButton_Left)) { + PickResult picked = + Pick(mouse_x, mouse_y, window, renderer, physics, vis_options); + if (picked.body >= 0) { + perturb.select = picked.body; + perturb.flexselect = picked.flex; + perturb.skinselect = picked.skin; + + // Compute the local position of the selected object in the world. + mjtNum tmp[3]; + mju_sub3(tmp, picked.point, data->xpos + 3 * picked.body); + mju_mulMatTVec(perturb.localpos, data->xmat + 9 * picked.body, tmp, 3, 3); + } else { + perturb.select = 0; + perturb.flexselect = -1; + perturb.skinselect = -1; + } + } + + // Right double click. + if (ImGui::IsMouseDoubleClicked(ImGuiMouseButton_Right)) { + PickResult picked = + Pick(mouse_x, mouse_y, window, renderer, physics, vis_options); + mju_copy3(camera.lookat, picked.point); + if (picked.body > 0 && io.KeyCtrl) { + camera.type = mjCAMERA_TRACKING; + camera.trackbodyid = picked.body; + camera.fixedcamid = -1; + camera_idx = 1; + } + } +} + +void HandleKeyboardEvents(Window* window, Renderer* renderer, + Physics* physics, mjvPerturb& perturb, + mjvOption& vis_options, mjvCamera& camera, + int& camera_idx) { + if (ImGui::GetIO().WantCaptureKeyboard) { + return; + } + + mjModel* model = physics->GetModel(); + + // Physics control shortcuts. + if (ImGui_IsChordJustPressed(ImGuiKey_Space)) { + physics->TogglePause(); + } else if (ImGui_IsChordJustPressed(ImGuiKey_Backspace)) { + physics->Reset(); + } + + // Camera shortcuts. + if (model) { + if (ImGui_IsChordJustPressed(ImGuiKey_Escape)) { + camera_idx = SetCamera(*model, camera, 0); + } else if (ImGui_IsChordJustPressed(ImGuiKey_LeftBracket)) { + camera_idx = SetCamera(*model, camera, camera_idx - 1); + } else if (ImGui_IsChordJustPressed(ImGuiKey_RightBracket)) { + camera_idx = SetCamera(*model, camera, camera_idx + 1); + } + } + + // Perturb shortcuts. + if (ImGui_IsChordJustPressed(ImGuiKey_PageUp)) { + SelectParentPerturb(model, perturb); + } + + // Visualization shortcuts. + if (ImGui_IsChordJustPressed(ImGuiKey_F6)) { + vis_options.frame = (vis_options.frame + 1) % mjNFRAME; + } else if (ImGui_IsChordJustPressed(ImGuiKey_F7)) { + vis_options.label = (vis_options.label + 1) % mjNLABEL; + } else if (ImGui_IsChordJustPressed(ImGuiKey_Comma)) { + ToggleFlag(vis_options.flags[mjVIS_ACTIVATION]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_Backslash)) { + ToggleFlag(vis_options.flags[mjVIS_MESHBVH]); + // } else if (ImGui_IsChordJustPressed(ImGuiKey_Backquote)) { + // ToggleFlag(vis_options.flags[mjVIS_BODYBVH]); + // } else if (ImGui_IsChordJustPressed(ImGuiKey_Quote)) { + // ToggleFlag(vis_options.flags[mjVIS_SCLINERTIA]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_Semicolon)) { + ToggleFlag(vis_options.flags[mjVIS_SKIN]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_U)) { + ToggleFlag(vis_options.flags[mjVIS_ACTUATOR]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_Q)) { + ToggleFlag(vis_options.flags[mjVIS_CAMERA]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_M)) { + ToggleFlag(vis_options.flags[mjVIS_COM]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_F)) { + ToggleFlag(vis_options.flags[mjVIS_CONTACTFORCE]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_C)) { + ToggleFlag(vis_options.flags[mjVIS_CONTACTPOINT]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_P)) { + ToggleFlag(vis_options.flags[mjVIS_CONTACTSPLIT]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_H)) { + ToggleFlag(vis_options.flags[mjVIS_CONVEXHULL]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_E)) { + ToggleFlag(vis_options.flags[mjVIS_CONSTRAINT]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_I)) { + ToggleFlag(vis_options.flags[mjVIS_ISLAND]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_J)) { + ToggleFlag(vis_options.flags[mjVIS_JOINT]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_Z)) { + ToggleFlag(vis_options.flags[mjVIS_LIGHT]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_B)) { + ToggleFlag(vis_options.flags[mjVIS_PERTFORCE]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_O)) { + ToggleFlag(vis_options.flags[mjVIS_PERTOBJ]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_Y)) { + ToggleFlag(vis_options.flags[mjVIS_RANGEFINDER]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_V)) { + ToggleFlag(vis_options.flags[mjVIS_TENDON]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_X)) { + ToggleFlag(vis_options.flags[mjVIS_TEXTURE]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_T)) { + ToggleFlag(vis_options.flags[mjVIS_TRANSPARENT]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_K)) { + ToggleFlag(vis_options.flags[mjVIS_AUTOCONNECT]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_G)) { + ToggleFlag(vis_options.flags[mjVIS_STATIC]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_0 | ImGuiMod_Shift)) { + ToggleFlag(vis_options.sitegroup[0]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_1 | ImGuiMod_Shift)) { + ToggleFlag(vis_options.sitegroup[1]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_2 | ImGuiMod_Shift)) { + ToggleFlag(vis_options.sitegroup[2]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_3 | ImGuiMod_Shift)) { + ToggleFlag(vis_options.sitegroup[3]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_4 | ImGuiMod_Shift)) { + ToggleFlag(vis_options.sitegroup[4]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_5 | ImGuiMod_Shift)) { + ToggleFlag(vis_options.sitegroup[5]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_0)) { + ToggleFlag(vis_options.geomgroup[0]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_1)) { + ToggleFlag(vis_options.geomgroup[1]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_2)) { + ToggleFlag(vis_options.geomgroup[2]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_3)) { + ToggleFlag(vis_options.geomgroup[3]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_4)) { + ToggleFlag(vis_options.geomgroup[4]); + } else if (ImGui_IsChordJustPressed(ImGuiKey_5)) { + ToggleFlag(vis_options.geomgroup[5]); + } +} + +} // namespace mujoco::toolbox diff --git a/src/experimental/toolbox/interaction.h b/src/experimental/toolbox/interaction.h new file mode 100644 index 00000000..ae623e0e --- /dev/null +++ b/src/experimental/toolbox/interaction.h @@ -0,0 +1,62 @@ +// Copyright 2025 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MUJOCO_SRC_EXPERIMENTAL_TOOLBOX_INTERACTION_H_ +#define MUJOCO_SRC_EXPERIMENTAL_TOOLBOX_INTERACTION_H_ + +#include "experimental/toolbox/physics.h" +#include "experimental/toolbox/renderer.h" +#include "experimental/toolbox/window.h" +#include + +namespace mujoco::toolbox { + +// The result of a pick operation. +struct PickResult { + mjtNum point[3]; // World coordinates + int body = -1; + int geom = -1; + int flex = -1; + int skin = -1; +}; + +// Returns information about the object (if any) under the mouse cursor. +PickResult Pick(float x, float y, Window* window, Renderer* renderer, + Physics* physics, const mjvOption& vis_options); + +// Updates the camera according to the requested index using this convention: +// +// 0 : selects the free camera (not defined in the model) +// 1 : selects the tracking camera (also not defined in the model) +// 2+ : selects a camera in the model; e.g. index 2 => model.cam[0]; +// +// The function returns the index of the used camera following the same +// convention. Note the returned index may differ from the request if the +// request was invalid (index was out of range or tracking camera was not +// available). +int SetCamera(const mjModel& model, mjvCamera& camera, int request_idx); + +// Handles canonical mouse events. +void HandleMouseEvents(Window* window, Renderer* renderer, Physics* physics, + mjvPerturb& perturb, mjvOption& vis_options, + mjvCamera& camera, int& camera_idx); + +// Handles canonical keyboard events. +void HandleKeyboardEvents(Window* window, Renderer* renderer, Physics* physics, + mjvPerturb& perturb, mjvOption& vis_options, + mjvCamera& camera, int& camera_idx); + +} // namespace mujoco::toolbox + +#endif // MUJOCO_SRC_EXPERIMENTAL_TOOLBOX_INTERACTION_H_ diff --git a/src/experimental/toolbox/physics.cc b/src/experimental/toolbox/physics.cc new file mode 100644 index 00000000..e1ad6e21 --- /dev/null +++ b/src/experimental/toolbox/physics.cc @@ -0,0 +1,222 @@ +// Copyright 2025 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "experimental/toolbox/physics.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "experimental/toolbox/helpers.h" +#include "experimental/toolbox/step_control.h" +#include + +namespace mujoco::toolbox { + +static mjtNum Timer() { + using Clock = std::chrono::steady_clock; + using Milliseconds = std::chrono::duration; + static Clock::time_point start = Clock::now(); + return Milliseconds(Clock::now() - start).count(); +} + +Physics::Physics(OnModelLoadedFn on_model_loaded) + : on_model_loaded_(std::move(on_model_loaded)) { + mjcb_time = Timer; +} + +Physics::~Physics() { Clear(); } + +void Physics::LoadModel(std::string model_file, const mjVFS* vfs) { + pending_load_ = std::move(model_file); + vfs_ = vfs; +} + +void Physics::ProcessPendingLoad() { + if (!pending_load_.has_value()) { + return; + } + + Clear(); + + std::string model_file = std::move(pending_load_.value()); + pending_load_.reset(); + + model_ = LoadMujocoModel(model_file, vfs_); + if (!model_) mju_error("Error loading model"); + + data_ = mj_makeData(model_); + if (!data_) mju_error("Error making data."); + + on_model_loaded_(model_file); + + InitHistory(); +} + +void Physics::Clear() { + if (model_) { + mj_deleteData(data_); + data_ = nullptr; + mj_deleteModel(model_); + model_ = nullptr; + + history_.clear(); + history_cursor_ = 0; + steps_ = 0; + GetStepControl().SetSpeed(100.f); + + error_ = ""; + } +} + +void Physics::Reset() { + mj_resetData(model_, data_); + mj_forward(model_, data_); + error_ = ""; + history_cursor_ = 0; +} + +bool Physics::Update(const mjvPerturb* perturb) { + ProcessPendingLoad(); + + if (!model_ || !data_) { + return false; + } + + if (data_) { + for (int i = 0; i < mjNTIMER; i++) { + data_->timer[i].duration = 0; + data_->timer[i].number = 0; + } + } + + if (!IsPaused()) { + mju_zero(data_->xfrc_applied, 6 * model_->nbody); + mjv_applyPerturbPose(model_, data_, perturb, 0); + mjv_applyPerturbForce(model_, data_, perturb); + } else { + mjv_applyPerturbPose(model_, data_, perturb, 1); + } + + if (IsPaused() && !single_step_) { + // run mj_forward, to update rendering and joint sliders + mj_forward(model_, data_); + if (pause_update_) { + mju_copy(data_->qacc_warmstart, data_->qacc, model_->nv); + } + + // When unpaused make sure we sync to immediately and step once. Without + // this we could step many times before rendering resulting in a noticeable + // delay before the simulation restarts (especially for large slowdowns) + GetStepControl().ForceSync(); + } else { + if (single_step_) { + GetStepControl().ForceSync(); + single_step_ = false; + } + + StepControl::Status status = GetStepControl().Advance(model_, data_); + if (status == StepControl::Status::kOk) { + AddToHistory(); + } else if (status == StepControl::Status::kAutoReset) { + Reset(); + } else if (status == StepControl::Status::kDiverged) { + for (mjtWarning w : StepControl::kDivergedWarnings) { + if (data_->warning[w].number > 0) { + paused_ = true; + error_ = mju_warningText(w, data_->warning[w].lastinfo); + } + } + } + } + + return true; +} + +bool Physics::UpdateState(mjtNum* state, unsigned int state_sig) { + ProcessPendingLoad(); + if (!model_ || !data_) { + return false; + } + mj_setState(model_, data_, state, state_sig); + mj_forward(model_, data_); + return true; +} + +void Physics::TogglePause() { paused_ = !paused_; } + +void Physics::RequestSingleStep() { single_step_ = true; } + +void Physics::InitHistory() { + const int state_size = mj_stateSize(model_, mjSTATE_INTEGRATION); + + // History buffer will be smaller of 2000 states or 100 MB. + constexpr int kMaxBytes = 1e8; + constexpr int kMaxHistory = 2000; + const int state_bytes = state_size * sizeof(mjtNum); + const int history_length = std::min(INT_MAX / state_bytes, kMaxHistory); + const int history_bytes = std::min(state_bytes * history_length, kMaxBytes); + const int num_history = history_bytes / state_bytes; + + history_.resize(num_history); + for (std::vector& state : history_) { + state.resize(state_size, 0); + } + history_cursor_ = 0; +} + +void Physics::AddToHistory() { + if (!history_.empty()) { + mjtNum* state = history_[history_cursor_].data(); + mj_getState(model_, data_, state, mjSTATE_INTEGRATION); + history_cursor_ = (history_cursor_ + 1) % history_.size(); + steps_++; + + // If we are adding to the history we didn't have a divergence error + error_ = ""; + } +} + +int Physics::LoadHistory(int offset) { + // No history to load. + if (steps_ == 0) { + return 0; + } + + // Pause simulation when entering history mode. + paused_ = true; + + // Ensure the offset is within a valid range. It's a negative value since + // we will be going backwards from the "latest" frame. + const int max_history = std::min(steps_, history_.size()); + offset = std::clamp(offset, -max_history + 1, 0); + + // Determine the index in the history buffer that corresponds to the frame + // index. + const int idx = (history_cursor_ + offset - 1) % history_.size(); + const mjtNum* state = history_[idx].data(); + + // Load the state into the data buffer. + mj_setState(model_, data_, state, mjSTATE_INTEGRATION); + mj_forward(model_, data_); + return offset; +} + +} // namespace mujoco::toolbox diff --git a/src/experimental/toolbox/physics.h b/src/experimental/toolbox/physics.h new file mode 100644 index 00000000..5cd960e4 --- /dev/null +++ b/src/experimental/toolbox/physics.h @@ -0,0 +1,126 @@ +// Copyright 2025 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MUJOCO_SRC_EXPERIMENTAL_TOOLBOX_PHYSICS_H_ +#define MUJOCO_SRC_EXPERIMENTAL_TOOLBOX_PHYSICS_H_ + +#include +#include +#include +#include +#include + + +#include +#include "experimental/toolbox/step_control.h" + +namespace mujoco::toolbox { + +// Owns the MuJoCo simulation state (e.g. mjModel and mjData) and is responsible +// for updating the state of the simulation. +class Physics { + public: + using OnModelLoadedFn = std::function; + + explicit Physics(OnModelLoadedFn on_model_loaded); + ~Physics(); + + Physics(const Physics&) = delete; + Physics& operator=(const Physics&) = delete; + + // Access the step controller + StepControl& GetStepControl() { return step_control_; } + + // Loads a model from the given path. An empty string will load an empty + // scene. + void LoadModel(std::string model_file, const mjVFS* vfs = nullptr); + + // Clears the simulation, clearing all loaded state. + void Clear(); + + // Resets the simulation using mj_resetData + void Reset(); + + // Advances the state of the simulation. + bool Update(const mjvPerturb* perturb); + + // Sets the state of the simulation. + bool UpdateState(mjtNum* state, unsigned int state_sig); + + // Renders the state of the simulation. + void Render(); + + // Returns true if the simulation is paused. + bool IsPaused() { return paused_; } + + // Pauses/unpauses the simulation. + void TogglePause(); + + // If the simulation is paused, will perform a single step on the next + // Update() call. + void RequestSingleStep(); + + // Returns the number of steps the simulation has taken. + int GetStepCount() const { return steps_; } + + // Returns the number of states in the history buffer. + int GetHistorySize() const { return history_.size(); } + + // Loads a state from the history buffer at the given offset into the current + // physics state. + // + // Calling this function will automatically pause the simulation. + int LoadHistory(int offset); + + // Selects the parent of the currently selected perturb object. + void SelectParentPerturb(); + + // Returns the MuJoCo data structures owned by this Simulation object. + mjModel* GetModel() { return model_; } + mjData* GetData() { return data_; } + + // Returns the error message from the simulation. + std::string_view GetError() const { return error_; } + + private: + void ProcessPendingLoad(); + void InitHistory(); + void AddToHistory(); + + mjModel* model_ = nullptr; + mjData* data_ = nullptr; + std::vector> history_; + int history_cursor_ = 0; + + OnModelLoadedFn on_model_loaded_; + + // If true and paused, d->qacc_warmstart is set to d->qacc after mj_forward + // which has the effect of making the constraint solver eventually converge + // while the simulation is paused. + bool pause_update_ = false; + + bool paused_ = false; + bool single_step_ = false; + int steps_ = 0; + std::optional pending_load_; + const mjVFS* vfs_; + + std::string error_; + + StepControl step_control_; +}; + +} // namespace mujoco::toolbox + +#endif // MUJOCO_SRC_EXPERIMENTAL_TOOLBOX_PHYSICS_H_ diff --git a/src/experimental/toolbox/renderer.cc b/src/experimental/toolbox/renderer.cc new file mode 100644 index 00000000..e4164ef8 --- /dev/null +++ b/src/experimental/toolbox/renderer.cc @@ -0,0 +1,82 @@ +// Copyright 2025 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "experimental/toolbox/renderer.h" + +#include +#include + +#include "experimental/toolbox/helpers.h" +#include + +namespace mujoco::toolbox { + +Renderer::Renderer(MakeContextFn make_context_fn) + : make_context_fn_(make_context_fn) { +} + +Renderer::~Renderer() { Deinit(); } + +void Renderer::Init(const mjModel* model) { + if (initialized_) { + Deinit(); + } + + if (model == nullptr) { + return; + } + + make_context_fn_(model, &render_context_); + mjv_defaultScene(&scene_); + mjv_makeScene(model, &scene_, 2000); + initialized_ = true; +} + +void Renderer::Deinit() { + mjv_freeScene(&scene_); + mjr_freeContext(&render_context_); + initialized_ = false; +} + +void Renderer::Sync(const mjModel* model, mjData* data, + const mjvPerturb* perturb, mjvCamera* camera, + const mjvOption* vis_option) { + mjv_updateScene(model, data, vis_option, perturb, camera, mjCAT_ALL, + &scene_); +} + +void Renderer::Render(const mjModel* model, mjData* data, + const mjvPerturb* perturb, mjvCamera* camera, + const mjvOption* vis_option, int width, int height) { + mjrRect main_viewport = {0, 0, width, height}; + mjr_render(main_viewport, data ? &scene_ : nullptr, &render_context_); + + auto now = std::chrono::steady_clock::now(); + auto delta_time = now - last_fps_update_; + const double interval = std::chrono::duration(delta_time).count(); + + ++frames_; + if (interval > 0.2) { // only update FPS stat at most 5 times per second + last_fps_update_ = now; + fps_ = frames_ / interval; + frames_ = 0; + } +} + +void Renderer::SaveScreenshot(const std::string& filename, int width, + int height) { + SaveScreenshotToWebp(width, height, &render_context_, filename); +} + +} // namespace mujoco::toolbox diff --git a/src/experimental/toolbox/renderer.h b/src/experimental/toolbox/renderer.h new file mode 100644 index 00000000..2891a45a --- /dev/null +++ b/src/experimental/toolbox/renderer.h @@ -0,0 +1,82 @@ +// Copyright 2025 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MUJOCO_SRC_EXPERIMENTAL_TOOLBOX_RENDERER_H_ +#define MUJOCO_SRC_EXPERIMENTAL_TOOLBOX_RENDERER_H_ + +#include +#include +#include +#include + +#include + +namespace mujoco::toolbox { + +// Renders the mujoco simulation and the imgui state into the active window +// using the filament rendering backend. +class Renderer { + public: + // Function that creates a mjrContext for the given model. We use a function + // to allow different mjrContext implementations to be created without + // requiring a direct dependency on them. + using MakeContextFn = std::function; + + explicit Renderer(MakeContextFn make_context_fn); + ~Renderer(); + + Renderer(const Renderer&) = delete; + Renderer& operator=(const Renderer&) = delete; + + // Initializes the renderer with the given mjModel. + void Init(const mjModel* model); + + // Updates the render scene with the current simulation state. + void Sync(const mjModel* model, mjData* data, const mjvPerturb* perturb, + mjvCamera* camera, const mjvOption* vis_option); + + // Renders the simulation state into the active window. Also renders the imgui + // state, but that is obtained directly from the ImGui library. + void Render(const mjModel* model, mjData* data, const mjvPerturb* perturb, + mjvCamera* camera, const mjvOption* vis_option, int width, + int height); + + // Saves a screenshot of the simulation state into the given file. + void SaveScreenshot(const std::string& filename, int width, int height); + + // Returns the mjvScene used by the renderer. + mjvScene& GetScene() { return scene_; } + + // Returns the current, average frame rate. + double GetFrameRate() const { return fps_; } + + private: + using TimePoint = std::chrono::time_point; + + // Resets the renderer; no rendering will occur until Init() is called again. + void Deinit(); + + MakeContextFn make_context_fn_; + mjrContext render_context_; + mjvScene scene_; + bool initialized_ = false; + + int frames_ = 0; + TimePoint last_fps_update_; + double fps_ = 0; +}; + +} // namespace mujoco::toolbox + +#endif // MUJOCO_SRC_EXPERIMENTAL_TOOLBOX_RENDERER_H_ diff --git a/src/experimental/toolbox/step_control.cc b/src/experimental/toolbox/step_control.cc new file mode 100644 index 00000000..311208b0 --- /dev/null +++ b/src/experimental/toolbox/step_control.cc @@ -0,0 +1,178 @@ +// Copyright 2025 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "experimental/toolbox/step_control.h" +#include +#include +#include +#include + +namespace mujoco::toolbox { + +float StepControl::GetSpeedMeasured() const { + return speed_measured_; +} + +float StepControl::GetSpeed() const { + return speed_; +} + +void StepControl::SetSpeed(float speed_percent_real_time) { + speed_ = std::clamp(speed_percent_real_time, .1f, 100.f); + ForceSync(); +} + +void StepControl::ForceSync() { force_sync_ = true; } + +void StepControl::GetNoiseParameters(float& ctrl_noise_scale, + float& ctrl_noise_rate) const { + ctrl_noise_scale = ctrl_noise_std_; + ctrl_noise_rate = ctrl_noise_rate_; +} + +void StepControl::SetNoiseParameters(float ctrl_noise_scale, + float ctrl_noise_rate) { + ctrl_noise_std_ = ctrl_noise_scale; + ctrl_noise_rate_ = ctrl_noise_rate; +} + +StepControl::Status StepControl::Advance(const mjModel* m, mjData* d) { + if (!m) { + return Status::kOk; + } + + const Clock::time_point start_cpu = Clock::now(); + const double slowdown = 100. / std::clamp(speed_, 0.001, 100.); + double elapsed_cpu = Seconds(start_cpu - sync_cpu_).count(); + double elapsed_sim = d->time - sync_sim_; + + bool resync = false; + + // Resync if we're forced to. + if (force_sync_) { + force_sync_ = false; + resync = true; + } + + // Resync if we've never synced. + if (sync_cpu_.time_since_epoch().count() == 0) { + resync = true; + } + + // Resync if any elapsed time is negative. + if (elapsed_cpu < 0 || elapsed_sim < 0) { + resync = true; + } + + // Resync if the distance from the target simulation time is bigger than + // sync_misalign_ (misalignment condition). + if (std::abs(elapsed_cpu / slowdown - elapsed_sim) > sync_misalign_) { + resync = true; + } + + if (resync) { + // Reset sync times. + sync_cpu_ = start_cpu; + sync_sim_ = d->time; + } + + // Stepping loop. + while (true) { + const Clock::time_point now_cpu = Clock::now(); + elapsed_cpu = Seconds(now_cpu - sync_cpu_).count(); + elapsed_sim = d->time - sync_sim_; + + // Stop stepping if simulation no longer lags cpu. + if (elapsed_sim * slowdown >= elapsed_cpu) { + return Status::kOk; + } + + // Stop stepping if simulation is taking too long to catch up. + // Note: 12ms == 70% of 1/60 seconds/frame. + constexpr Clock::duration kMaxCpuTimeForSim = std::chrono::milliseconds(12); + if (now_cpu - start_cpu >= kMaxCpuTimeForSim) { + // Note: GetSpeed() and GetSpeedMeasured() will be different in this case. + return Status::kOk; + } + + // Measure slowdown here in first viable in-sync step. This update location + // is chosen to minimize visual noise caused by changing measurements. + if (elapsed_sim > 0) { + double measured_slowdown = elapsed_cpu / elapsed_sim; + speed_measured_ = 100. / measured_slowdown; + } + + mjtNum prev_time = d->time; + InjectNoise(m, d); + mj_step(m, d); + + if (mjDISABLED(mjDSBL_AUTORESET)) { + for (mjtWarning w : kDivergedWarnings) { + // Stop stepping if the simulation diverged. + if (d->warning[w].number > 0) { + return Status::kDiverged; + } + } + } else { + // Stop stepping if we auto reset. + if (d->time < prev_time) { + return Status::kAutoReset; + } + } + + // Stop after one step if we resynced; next iteration will deal with timing. + if (resync) { + return Status::kOk; + } + } + + return Status::kDiverged; // Unreachable +} + +void StepControl::InjectNoise(const mjModel* m, mjData* d) { + // no noise, return + if (ctrl_noise_std_ <= 0) { + return; + } + + // convert rate and scale to discrete time (Ornstein–Uhlenbeck) + mjtNum rate = mju_exp(-m->opt.timestep / ctrl_noise_rate_); + mjtNum scale = ctrl_noise_std_ * mju_sqrt(1-rate*rate); + + for (int i = 0; i < m->nu; i++) { + mjtNum bottom = 0; + mjtNum top = 0; + mjtNum midpoint = 0; + mjtNum halfrange = 1; + if (m->actuator_ctrllimited[i]) { + bottom = m->actuator_ctrlrange[2*i]; + top = m->actuator_ctrlrange[2*i+1]; + midpoint = 0.5 * (top + bottom); // target of exponential decay + halfrange = 0.5 * (top - bottom); // scales noise + } + + // exponential convergence to midpoint at ctrl_noise_rate + d->ctrl[i] = rate * d->ctrl[i] + (1-rate) * midpoint; + + // add noise + d->ctrl[i] += scale * halfrange * mju_standardNormal(nullptr); + + // clip to range if limited + if (m->actuator_ctrllimited[i]) { + d->ctrl[i] = mju_clip(d->ctrl[i], bottom, top); + } + } +} + +} // namespace mujoco::toolbox diff --git a/src/experimental/toolbox/step_control.h b/src/experimental/toolbox/step_control.h new file mode 100644 index 00000000..25484b96 --- /dev/null +++ b/src/experimental/toolbox/step_control.h @@ -0,0 +1,93 @@ +// Copyright 2025 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MUJOCO_SRC_EXPERIMENTAL_TOOLBOX_STEP_CONTROL_H_ +#define MUJOCO_SRC_EXPERIMENTAL_TOOLBOX_STEP_CONTROL_H_ + +#include +#include + +#include + +namespace mujoco::toolbox { + +using Seconds = std::chrono::duration; +using Clock = std::chrono::steady_clock; + +// State and logic for physics synchronization and stepping. +class StepControl { + public: + enum class Status { + kOk, + + // Simulation diverged with autoreset enabled. + kAutoReset, + + // Simulation diverged with autoreset disabled. + // Note: Consider reporting mjData warning diagnostics in kDivergedWarnings. + kDiverged, + }; + + // List of warnings that are checked to determine simulation divergence. + static constexpr mjtWarning kDivergedWarnings[] = { + mjWARN_BADQACC, mjWARN_BADQVEL, mjWARN_BADQPOS}; + + // Steps physics forward, respecting speed settings and refresh budget. + Status Advance(const mjModel* m, mjData* d); + + // Ensures the next call to Advance() will synchronize time and step once. + void ForceSync(); + + // Gets/sets the desired simulation speed as a percentage of real time. + float GetSpeed() const; + float GetSpeedMeasured() const; + void SetSpeed(float speed); // speed is clamped to [0.1%, 100%] + + // Gets/sets the control noise parameters applied before stepping. + void GetNoiseParameters(float& noise_scale, float& noise_rate) const; + void SetNoiseParameters(float noise_scale, float noise_rate); + + private: + std::string AdvanceOneStep(const mjModel* m, mjData* d); + + void InjectNoise(const mjModel* m, mjData* d); + + // Control noise standard deviation + double ctrl_noise_std_ = 0; + + // Control noise correlation rate + double ctrl_noise_rate_ = 0; + + // Desired simulation speed as a percentage of real time + float speed_ = 100; + + // Measured simulation speed as a percentage of real time + float speed_measured_ = -1; + + // If true, the next call to Advance() will synchronize time step once. + bool force_sync_ = true; + + // CPU time (aka wall time) of the last synchronization event + std::chrono::time_point sync_cpu_; + + // Simulation time of the last synchronization event + mjtNum sync_sim_ = 0; + + // Maximum mis-alignment before re-sync (simulation seconds) + double sync_misalign_ = .1; +}; + +} // namespace mujoco::toolbox + +#endif // MUJOCO_SRC_EXPERIMENTAL_TOOLBOX_STEP_CONTROL_H_ diff --git a/src/experimental/toolbox/window.cc b/src/experimental/toolbox/window.cc new file mode 100644 index 00000000..8afe5fba --- /dev/null +++ b/src/experimental/toolbox/window.cc @@ -0,0 +1,168 @@ +// Copyright 2025 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "experimental/toolbox/window.h" + +#include +#include + +#include "third_party/SDL2/include/SDL.h" +#include "third_party/SDL2/include/SDL_error.h" +#include "third_party/SDL2/include/SDL_events.h" +#include "third_party/SDL2/include/SDL_hints.h" +#include "third_party/SDL2/include/SDL_syswm.h" +#include "third_party/SDL2/include/SDL_version.h" +#include "third_party/SDL2/include/SDL_video.h" +#include "third_party/dear_imgui/backends/imgui_impl_sdl2.h" +#include "third_party/dear_imgui/imgui.h" +#include "third_party/implot/implot.h" +#include "experimental/toolbox/helpers.h" +#include + +namespace mujoco::toolbox { + +static void InitImGui(SDL_Window* window, const LoadAssetFn& load_asset_fn) { + ImGui::CreateContext(); + ImPlot::CreateContext(); + + ImGuiIO& io = ImGui::GetIO(); + io.ConfigFlags |= ImGuiConfigFlags_DockingEnable; + io.IniFilename = nullptr; + ImGui::StyleColorsDark(); + ImGui_ImplSDL2_InitForOther(window); + +#ifndef __EMSCRIPTEN__ + // TODO: Get font loading working for wasm. + ImFontConfig main_cfg; + main_cfg.OversampleH = 8; + main_cfg.OversampleV = 4; + main_cfg.GlyphExtraSpacing.x = 0.3f; + + auto main_font = load_asset_fn("Roboto-Regular.ttf"); + io.Fonts->AddFontFromMemoryTTF(main_font.data(), main_font.size(), 18.f, + &main_cfg); + + ImFontConfig icon_cfg; + icon_cfg.OversampleH = 3; + icon_cfg.OversampleV = 3; + icon_cfg.MergeMode = true; + icon_cfg.GlyphMinAdvanceX = 14.0f; + auto icon_font = load_asset_fn("fontawesome-webfont.ttf"); + constexpr ImWchar icon_ranges[] = {0xf000, 0xf3ff, 0x000}; + io.Fonts->AddFontFromMemoryTTF(icon_font.data(), icon_font.size(), 14.f, + &icon_cfg, icon_ranges); +#endif + io.Fonts->Build(); +} + +Window::Window(std::string_view title, int width, int height, Config config, + const LoadAssetFn& load_asset_fn) + : width_(width), height_(height), config_(config) { + SDL_SetHint(SDL_HINT_FRAMEBUFFER_ACCELERATION, "1"); + SDL_GL_SetAttribute(SDL_GL_MULTISAMPLEBUFFERS, 1); + SDL_GL_SetAttribute(SDL_GL_MULTISAMPLESAMPLES, 16); + if (SDL_Init(SDL_INIT_VIDEO | SDL_INIT_EVENTS) != 0) { + mju_error("Error initializing SDL: %s", SDL_GetError()); + } + + int window_flags = SDL_WINDOW_RESIZABLE | SDL_WINDOW_ALLOW_HIGHDPI; + + if (config == kFilamentVulkan) { + window_flags |= SDL_WINDOW_VULKAN; + } else if (config == kFilamentWebGL) { + window_flags |= SDL_WINDOW_OPENGL; + SDL_GL_SetAttribute(SDL_GL_CONTEXT_MAJOR_VERSION, 3); + SDL_GL_SetAttribute(SDL_GL_CONTEXT_MINOR_VERSION, 0); + SDL_GL_SetAttribute(SDL_GL_CONTEXT_PROFILE_MASK, SDL_GL_CONTEXT_PROFILE_ES); + } else if (config == kMujocoOpenGL || config == kFilamentOpenGL) { + window_flags |= SDL_WINDOW_OPENGL; + SDL_GL_SetAttribute(SDL_GL_CONTEXT_MAJOR_VERSION, 3); + SDL_GL_SetAttribute(SDL_GL_CONTEXT_MINOR_VERSION, 0); + } else { + mju_error("Unsupported window config: %d", config); + } + + sdl_window_ = + SDL_CreateWindow(title.data(), SDL_WINDOWPOS_UNDEFINED, + SDL_WINDOWPOS_UNDEFINED, width, height, window_flags); + if (!sdl_window_) { + mju_error("Error creating window: %s", SDL_GetError()); + } + + InitImGui(sdl_window_, load_asset_fn); + + if (config == kFilamentWebGL || config == kMujocoOpenGL) { + SDL_GLContext gl_context = SDL_GL_CreateContext(sdl_window_); + SDL_GL_MakeCurrent(sdl_window_, gl_context); + } + + #ifdef __linux__ + SDL_SysWMinfo wmi; + SDL_VERSION(&wmi.version); + SDL_GetWindowWMInfo(sdl_window_, &wmi); + native_window_ = reinterpret_cast(wmi.info.x11.window); + #endif +} + +Window::~Window() { + SDL_DestroyWindow(sdl_window_); + SDL_Quit(); +} + +void Window::SetTitle(std::string_view title) { + SDL_SetWindowTitle(sdl_window_, title.data()); +} + +std::string Window::GetDropFile() { + std::string tmp; + std::swap(tmp, drop_file_); + return tmp; +} + +Window::Status Window::ProcessEvents() { + SDL_Event event; + while (SDL_PollEvent(&event)) { + ImGui_ImplSDL2_ProcessEvent(&event); + + if (event.type == SDL_QUIT) { + should_exit_ = true; + } else if (event.type == SDL_APP_WILLENTERBACKGROUND) { + should_exit_ = true; + } else if (event.type == SDL_WINDOWEVENT) { + if (event.window.event == SDL_WINDOWEVENT_RESIZED) { + SDL_GetWindowSize(sdl_window_, &width_, &height_); + int drawable_width = width_; + int drawable_height = height_; + SDL_GL_GetDrawableSize(sdl_window_, &drawable_width, &drawable_height); + scale_ = (float)drawable_width / (float)width_; + } + } else if (event.type == SDL_DROPFILE) { + drop_file_ = event.drop.file; + } + } + + ImGui_ImplSDL2_NewFrame(); + ImGui::NewFrame(); + + return should_exit_ ? kQuitting : kRunning; +} + +void Window::Present() { + // Filament (with the exception of WebGL) handles the swapchain internally. + if (config_ != kFilamentVulkan && config_ != kFilamentOpenGL) { + SDL_GL_SwapWindow(sdl_window_); + } +} + +} // namespace mujoco::toolbox diff --git a/src/experimental/toolbox/window.h b/src/experimental/toolbox/window.h new file mode 100644 index 00000000..8e78ce24 --- /dev/null +++ b/src/experimental/toolbox/window.h @@ -0,0 +1,87 @@ +// Copyright 2025 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MUJOCO_SRC_EXPERIMENTAL_TOOLBOX_WINDOW_H_ +#define MUJOCO_SRC_EXPERIMENTAL_TOOLBOX_WINDOW_H_ + +#include +#include + +#include "experimental/toolbox/helpers.h" +#include "third_party/SDL2/include/SDL_video.h" + +namespace mujoco::toolbox { + +// Platform-independent window abstraction using SDL2. +// +// Initializes the SDL2 and ImGui libraries, creates/owns the native window, and +// handles events from the window. +class Window { + public: + // Configures the window for the specified rendering backend. + enum Config { + kMujocoOpenGL, + kFilamentVulkan, + kFilamentOpenGL, + kFilamentWebGL, + }; + + Window(std::string_view title, int width, int height, Config config, + const LoadAssetFn& load_asset_fn); + ~Window(); + + Window(const Window&) = delete; + Window& operator=(const Window&) = delete; + + // The status of the window. + enum Status { + kRunning, + kQuitting, + }; + + // Processes all pendings window events, returning the status of the window. + Status ProcessEvents(); + + // Swaps and presents the window buffer. + void Present(); + + // Sets the title of the window. + void SetTitle(std::string_view title); + + // Returns the current size of the window. + int GetWidth() const { return width_; } + int GetHeight() const { return height_; } + float GetScale() const { return scale_; } + + // Returns the path to a file that was dropped on the window. Once called, + // the value will be cleared until the next time a file is dropped. + std::string GetDropFile(); + + // Returns the handle to the underlying native window. + void* GetNativeWindowHandle() { return native_window_; } + + private: + int width_ = 0; + int height_ = 0; + float scale_ = 1.0f; + Config config_; + void* native_window_ = nullptr; + SDL_Window* sdl_window_ = nullptr; + bool should_exit_ = false; + std::string drop_file_; +}; + +} // namespace mujoco::toolbox + +#endif // MUJOCO_SRC_EXPERIMENTAL_TOOLBOX_WINDOW_H_