Configure Copybara export for Dear ImGui and ImPlot Python bindings
Following the export declarations in Dear ImGui and ImPlot METADATA, this change updates MuJoCo's Copybara configuration (copy.bara.sky) to export and transform the Python bindings. `//third_party/dear_imgui/google/py` exports to `python/mujoco/experimental/dear_imgui` and `//third_party/implot/google/py` exports to `python/mujoco/experimental/implot`. PiperOrigin-RevId: 925293624 Change-Id: Ie6e32d247a6f7fc24bb36ae7060f2075d8efeb26
This commit is contained in:
Matija Kecman
committed by
Copybara-Service
Copybara-Service
parent
062b0f1ea6
commit
4cf4a5665d
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,326 @@
|
||||
// Copyright 2026 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
|
||||
//
|
||||
// https://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_PYTHON_EXPERIMENTAL_DEAR_IMGUI_DEAR_IMGUI_MACROS_H_
|
||||
#define MUJOCO_PYTHON_EXPERIMENTAL_DEAR_IMGUI_DEAR_IMGUI_MACROS_H_
|
||||
|
||||
// WARNING: This file is intended for internal use by dear_imgui libraries ONLY!
|
||||
//
|
||||
// The macros defined here use short, generic names (DEF0, ARG_ID, etc.) and
|
||||
// are NOT #undef'd. Including this header elsewhere may cause naming conflicts.
|
||||
//
|
||||
// Define NAMESPACE to be the ImGui library you're binding before including this
|
||||
// header, e.g. #define NAMESPACE ImGui
|
||||
//
|
||||
// ============================================================================
|
||||
// Quick Reference
|
||||
// ============================================================================
|
||||
//
|
||||
// DEFn(Name, Args...) // Binds NAMESPACE::Name as Name in Python
|
||||
// DEFn_AS(CppName, PyName, Args...) // Binds NAMESPACE::CppName as PyName
|
||||
// DEFn_F(PyName, Args..., { CppBody }) // Custom implementation
|
||||
//
|
||||
// Where 'n' is the number of arguments (0-9).
|
||||
//
|
||||
// ============================================================================
|
||||
// Argument Format
|
||||
// ============================================================================
|
||||
//
|
||||
// Each argument is a tuple: (Type, name, DefaultValue)
|
||||
//
|
||||
// - No default value: (ImString, label, ) // trailing comma required
|
||||
// - With default value: (int, flags, = 0) // include the '='
|
||||
// - Complex defaults: (const ImVec2&, size, = ImVec2_Zero)
|
||||
//
|
||||
// NOTE: Default values cannot contain commas. Use predefined constants like
|
||||
// ImVec2_Zero, ImVec4_One, etc.
|
||||
//
|
||||
// ============================================================================
|
||||
// Examples (from dear_imgui.cc)
|
||||
// ============================================================================
|
||||
//
|
||||
// Simple binding - NAMESPACE::End() exposed as End():
|
||||
// DEF0(End);
|
||||
//
|
||||
// Binding with arguments:
|
||||
// DEF4(BeginChild,
|
||||
// (ImString, str_id, ),
|
||||
// (const ImVec2&, size, = ImVec2_Zero),
|
||||
// (ImGuiChildFlags, child_flags, = 0),
|
||||
// (ImGuiWindowFlags, window_flags, = 0));
|
||||
//
|
||||
// Overloaded function - NAMESPACE::BeginChild(ImGuiID) exposed as BeginChildId():
|
||||
// DEF4_AS(BeginChild, BeginChildId,
|
||||
// (ImGuiID, id, ),
|
||||
// (const ImVec2&, size, = ImVec2_Zero),
|
||||
// (ImGuiChildFlags, child_flags, = 0),
|
||||
// (ImGuiWindowFlags, window_flags, = 0));
|
||||
//
|
||||
// Custom implementation using DEFn_F is needed when:
|
||||
//
|
||||
// 1. Variadic functions (e.g., Text, TextColored)
|
||||
//
|
||||
// C++ variadic functions (those with "...") cannot be bound directly
|
||||
// because the type/count of arguments is unknown at compile time. Use a
|
||||
// wrapper that calls the function with a fixed format. Also note that
|
||||
// user-controlled format strings are a security risk (format string
|
||||
// attacks). Always use "%s":
|
||||
//
|
||||
// DEF1_F(Text, (ImString, txt, ), {
|
||||
// return NAMESPACE::Text("%s", txt);
|
||||
// });
|
||||
//
|
||||
// 2. Output pointer parameters (e.g., Checkbox, SliderFloat)
|
||||
//
|
||||
// Python doesn't have output pointers, so return modified values as a
|
||||
// tuple:
|
||||
//
|
||||
// DEF2_F(Checkbox, (ImString, label, ), (bool*, v, ), {
|
||||
// auto result = NAMESPACE::Checkbox(label, v);
|
||||
// return std::make_tuple(result, *v);
|
||||
// });
|
||||
//
|
||||
// 3. Type conversions (e.g., Image, ImageButton)
|
||||
//
|
||||
// Some C++ types don't have Python equivalents. For example, ImTextureID
|
||||
// is a void* (opaque pointer), which pybind11 can't automatically convert.
|
||||
// Accept a Python-friendly type (like long) and cast it:
|
||||
//
|
||||
// DEF2_F(Image, (long, tex_id, ), (const ImVec2&, size, ), {
|
||||
// return NAMESPACE::Image(reinterpret_cast<void*>(tex_id), size);
|
||||
// });
|
||||
|
||||
// ============================================================================
|
||||
// Internal helper macros (not intended to be called directly by binding code)
|
||||
// ============================================================================
|
||||
|
||||
// Extracts the type and name of an argument tuple.
|
||||
// Example: ARG_DECL((float, alpha, = 1.0f)) -> float alpha
|
||||
#define ARG_DECL_X(T_, N_, V_) T_ N_
|
||||
#define ARG_DECL(A_) ARG_DECL_X A_
|
||||
|
||||
// Extracts the identifier of an argument tuple.
|
||||
// Example: ARG_ID((float, alpha, = 1.0f)) -> alpha
|
||||
#define ARG_ID_X(T_, N_, V_) N_
|
||||
#define ARG_ID(A_) ARG_ID_X A_
|
||||
|
||||
// Extracts the name of an argument tuple as a quoted string literal.
|
||||
// Example: ARG_NAME((float, alpha, = 1.0f)) -> "alpha"
|
||||
#define ARG_NAME_X(T_, N_, V_) #N_
|
||||
#define ARG_NAME(A_) ARG_NAME_X A_
|
||||
|
||||
// Extracts the default value of an argument tuple.
|
||||
// Example: ARG_DEFVAL((float, alpha, = 1.0f)) -> = 1.0f
|
||||
#define ARG_DEFVAL_X(T_, N_, V_) V_
|
||||
#define ARG_DEFVAL(A_) ARG_DEFVAL_X A_
|
||||
|
||||
// ============================================================================
|
||||
// Public macros for binding code
|
||||
// ============================================================================
|
||||
|
||||
//
|
||||
#define DEF0_F(N, FN) \
|
||||
m.def(#N, []( \
|
||||
) FN \
|
||||
);
|
||||
|
||||
#define DEF1_F(N, A1, FN) \
|
||||
m.def(#N, []( \
|
||||
ARG_DECL(A1) \
|
||||
) FN, \
|
||||
py::arg(ARG_NAME(A1)) ARG_DEFVAL(A1) \
|
||||
);
|
||||
|
||||
#define DEF2_F(N, A1, A2, FN) \
|
||||
m.def(#N, []( \
|
||||
ARG_DECL(A1), \
|
||||
ARG_DECL(A2) \
|
||||
) FN, \
|
||||
py::arg(ARG_NAME(A1)) ARG_DEFVAL(A1), \
|
||||
py::arg(ARG_NAME(A2)) ARG_DEFVAL(A2) \
|
||||
);
|
||||
|
||||
#define DEF3_F(N, A1, A2, A3, FN) \
|
||||
m.def(#N, []( \
|
||||
ARG_DECL(A1), \
|
||||
ARG_DECL(A2), \
|
||||
ARG_DECL(A3) \
|
||||
) FN, \
|
||||
py::arg(ARG_NAME(A1)) ARG_DEFVAL(A1), \
|
||||
py::arg(ARG_NAME(A2)) ARG_DEFVAL(A2), \
|
||||
py::arg(ARG_NAME(A3)) ARG_DEFVAL(A3) \
|
||||
);
|
||||
|
||||
#define DEF4_F(N, A1, A2, A3, A4, FN) \
|
||||
m.def(#N, []( \
|
||||
ARG_DECL(A1), \
|
||||
ARG_DECL(A2), \
|
||||
ARG_DECL(A3), \
|
||||
ARG_DECL(A4) \
|
||||
) FN, \
|
||||
py::arg(ARG_NAME(A1)) ARG_DEFVAL(A1), \
|
||||
py::arg(ARG_NAME(A2)) ARG_DEFVAL(A2), \
|
||||
py::arg(ARG_NAME(A3)) ARG_DEFVAL(A3), \
|
||||
py::arg(ARG_NAME(A4)) ARG_DEFVAL(A4) \
|
||||
);
|
||||
|
||||
#define DEF5_F(N, A1, A2, A3, A4, A5, FN) \
|
||||
m.def(#N, []( \
|
||||
ARG_DECL(A1), \
|
||||
ARG_DECL(A2), \
|
||||
ARG_DECL(A3), \
|
||||
ARG_DECL(A4), \
|
||||
ARG_DECL(A5) \
|
||||
) FN, \
|
||||
py::arg(ARG_NAME(A1)) ARG_DEFVAL(A1), \
|
||||
py::arg(ARG_NAME(A2)) ARG_DEFVAL(A2), \
|
||||
py::arg(ARG_NAME(A3)) ARG_DEFVAL(A3), \
|
||||
py::arg(ARG_NAME(A4)) ARG_DEFVAL(A4), \
|
||||
py::arg(ARG_NAME(A5)) ARG_DEFVAL(A5) \
|
||||
);
|
||||
|
||||
#define DEF6_F(N, A1, A2, A3, A4, A5, A6, FN) \
|
||||
m.def(#N, []( \
|
||||
ARG_DECL(A1), \
|
||||
ARG_DECL(A2), \
|
||||
ARG_DECL(A3), \
|
||||
ARG_DECL(A4), \
|
||||
ARG_DECL(A5), \
|
||||
ARG_DECL(A6) \
|
||||
) FN, \
|
||||
py::arg(ARG_NAME(A1)) ARG_DEFVAL(A1), \
|
||||
py::arg(ARG_NAME(A2)) ARG_DEFVAL(A2), \
|
||||
py::arg(ARG_NAME(A3)) ARG_DEFVAL(A3), \
|
||||
py::arg(ARG_NAME(A4)) ARG_DEFVAL(A4), \
|
||||
py::arg(ARG_NAME(A5)) ARG_DEFVAL(A5), \
|
||||
py::arg(ARG_NAME(A6)) ARG_DEFVAL(A6) \
|
||||
);
|
||||
|
||||
#define DEF7_F(N, A1, A2, A3, A4, A5, A6, A7, FN) \
|
||||
m.def(#N, []( \
|
||||
ARG_DECL(A1), \
|
||||
ARG_DECL(A2), \
|
||||
ARG_DECL(A3), \
|
||||
ARG_DECL(A4), \
|
||||
ARG_DECL(A5), \
|
||||
ARG_DECL(A6), \
|
||||
ARG_DECL(A7) \
|
||||
) FN, \
|
||||
py::arg(ARG_NAME(A1)) ARG_DEFVAL(A1), \
|
||||
py::arg(ARG_NAME(A2)) ARG_DEFVAL(A2), \
|
||||
py::arg(ARG_NAME(A3)) ARG_DEFVAL(A3), \
|
||||
py::arg(ARG_NAME(A4)) ARG_DEFVAL(A4), \
|
||||
py::arg(ARG_NAME(A5)) ARG_DEFVAL(A5), \
|
||||
py::arg(ARG_NAME(A6)) ARG_DEFVAL(A6), \
|
||||
py::arg(ARG_NAME(A7)) ARG_DEFVAL(A7) \
|
||||
);
|
||||
|
||||
#define DEF8_F(N, A1, A2, A3, A4, A5, A6, A7, A8, FN) \
|
||||
m.def(#N, []( \
|
||||
ARG_DECL(A1), \
|
||||
ARG_DECL(A2), \
|
||||
ARG_DECL(A3), \
|
||||
ARG_DECL(A4), \
|
||||
ARG_DECL(A5), \
|
||||
ARG_DECL(A6), \
|
||||
ARG_DECL(A7), \
|
||||
ARG_DECL(A8) \
|
||||
) FN, \
|
||||
py::arg(ARG_NAME(A1)) ARG_DEFVAL(A1), \
|
||||
py::arg(ARG_NAME(A2)) ARG_DEFVAL(A2), \
|
||||
py::arg(ARG_NAME(A3)) ARG_DEFVAL(A3), \
|
||||
py::arg(ARG_NAME(A4)) ARG_DEFVAL(A4), \
|
||||
py::arg(ARG_NAME(A5)) ARG_DEFVAL(A5), \
|
||||
py::arg(ARG_NAME(A6)) ARG_DEFVAL(A6), \
|
||||
py::arg(ARG_NAME(A7)) ARG_DEFVAL(A7), \
|
||||
py::arg(ARG_NAME(A8)) ARG_DEFVAL(A8) \
|
||||
);
|
||||
|
||||
#define DEF9_F(N, A1, A2, A3, A4, A5, A6, A7, A8, A9, FN) \
|
||||
m.def(#N, []( \
|
||||
ARG_DECL(A1), \
|
||||
ARG_DECL(A2), \
|
||||
ARG_DECL(A3), \
|
||||
ARG_DECL(A4), \
|
||||
ARG_DECL(A5), \
|
||||
ARG_DECL(A6), \
|
||||
ARG_DECL(A7), \
|
||||
ARG_DECL(A8), \
|
||||
ARG_DECL(A9) \
|
||||
) FN, \
|
||||
py::arg(ARG_NAME(A1)) ARG_DEFVAL(A1), \
|
||||
py::arg(ARG_NAME(A2)) ARG_DEFVAL(A2), \
|
||||
py::arg(ARG_NAME(A3)) ARG_DEFVAL(A3), \
|
||||
py::arg(ARG_NAME(A4)) ARG_DEFVAL(A4), \
|
||||
py::arg(ARG_NAME(A5)) ARG_DEFVAL(A5), \
|
||||
py::arg(ARG_NAME(A6)) ARG_DEFVAL(A6), \
|
||||
py::arg(ARG_NAME(A7)) ARG_DEFVAL(A7), \
|
||||
py::arg(ARG_NAME(A8)) ARG_DEFVAL(A8), \
|
||||
py::arg(ARG_NAME(A9)) ARG_DEFVAL(A9) \
|
||||
);
|
||||
|
||||
#define DEF10_F(N, A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, FN) \
|
||||
m.def(#N, []( \
|
||||
ARG_DECL(A1), \
|
||||
ARG_DECL(A2), \
|
||||
ARG_DECL(A3), \
|
||||
ARG_DECL(A4), \
|
||||
ARG_DECL(A5), \
|
||||
ARG_DECL(A6), \
|
||||
ARG_DECL(A7), \
|
||||
ARG_DECL(A8), \
|
||||
ARG_DECL(A9), \
|
||||
ARG_DECL(A10) \
|
||||
) FN, \
|
||||
py::arg(ARG_NAME(A1)) ARG_DEFVAL(A1), \
|
||||
py::arg(ARG_NAME(A2)) ARG_DEFVAL(A2), \
|
||||
py::arg(ARG_NAME(A3)) ARG_DEFVAL(A3), \
|
||||
py::arg(ARG_NAME(A4)) ARG_DEFVAL(A4), \
|
||||
py::arg(ARG_NAME(A5)) ARG_DEFVAL(A5), \
|
||||
py::arg(ARG_NAME(A6)) ARG_DEFVAL(A6), \
|
||||
py::arg(ARG_NAME(A7)) ARG_DEFVAL(A7), \
|
||||
py::arg(ARG_NAME(A8)) ARG_DEFVAL(A8), \
|
||||
py::arg(ARG_NAME(A9)) ARG_DEFVAL(A9), \
|
||||
py::arg(ARG_NAME(A10)) ARG_DEFVAL(A10) \
|
||||
);
|
||||
|
||||
// NOLINTBEGIN(whitespace/line_length)
|
||||
|
||||
#define DEF0_AS(N, AS) DEF0_F(AS, { return NAMESPACE::N(); } )
|
||||
#define DEF1_AS(N, AS, A1) DEF1_F(AS, A1, { return NAMESPACE::N(ARG_ID(A1)); } )
|
||||
#define DEF2_AS(N, AS, A1, A2) DEF2_F(AS, A1, A2, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2)); } )
|
||||
#define DEF3_AS(N, AS, A1, A2, A3) DEF3_F(AS, A1, A2, A3, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3)); } )
|
||||
#define DEF4_AS(N, AS, A1, A2, A3, A4) DEF4_F(AS, A1, A2, A3, A4, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3), ARG_ID(A4)); } )
|
||||
#define DEF5_AS(N, AS, A1, A2, A3, A4, A5) DEF5_F(AS, A1, A2, A3, A4, A5, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3), ARG_ID(A4), ARG_ID(A5)); } )
|
||||
#define DEF6_AS(N, AS, A1, A2, A3, A4, A5, A6) DEF6_F(AS, A1, A2, A3, A4, A5, A6, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3), ARG_ID(A4), ARG_ID(A5), ARG_ID(A6)); } )
|
||||
#define DEF7_AS(N, AS, A1, A2, A3, A4, A5, A6, A7) DEF7_F(AS, A1, A2, A3, A4, A5, A6, A7, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3), ARG_ID(A4), ARG_ID(A5), ARG_ID(A6), ARG_ID(A7)); } )
|
||||
#define DEF8_AS(N, AS, A1, A2, A3, A4, A5, A6, A7, A8) DEF8_F(AS, A1, A2, A3, A4, A5, A6, A7, A8, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3), ARG_ID(A4), ARG_ID(A5), ARG_ID(A6), ARG_ID(A7), ARG_ID(A8)); } )
|
||||
#define DEF9_AS(N, AS, A1, A2, A3, A4, A5, A6, A7, A8, A9) DEF9_F(AS, A1, A2, A3, A4, A5, A6, A7, A8, A9, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3), ARG_ID(A4), ARG_ID(A5), ARG_ID(A6), ARG_ID(A7), ARG_ID(A8), ARG_ID(A9)); } )
|
||||
#define DEF10_AS(N, AS, A1, A2, A3, A4, A5, A6, A7, A8, A9, A10) DEF10_F(AS, A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3), ARG_ID(A4), ARG_ID(A5), ARG_ID(A6), ARG_ID(A7), ARG_ID(A8), ARG_ID(A9), ARG_ID(A10)); } )
|
||||
|
||||
#define DEF0(N) DEF0_F(N, { return NAMESPACE::N(); } )
|
||||
#define DEF1(N, A1) DEF1_F(N, A1, { return NAMESPACE::N(ARG_ID(A1)); } )
|
||||
#define DEF2(N, A1, A2) DEF2_F(N, A1, A2, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2)); } )
|
||||
#define DEF3(N, A1, A2, A3) DEF3_F(N, A1, A2, A3, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3)); } )
|
||||
#define DEF4(N, A1, A2, A3, A4) DEF4_F(N, A1, A2, A3, A4, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3), ARG_ID(A4)); } )
|
||||
#define DEF5(N, A1, A2, A3, A4, A5) DEF5_F(N, A1, A2, A3, A4, A5, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3), ARG_ID(A4), ARG_ID(A5)); } )
|
||||
#define DEF6(N, A1, A2, A3, A4, A5, A6) DEF6_F(N, A1, A2, A3, A4, A5, A6, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3), ARG_ID(A4), ARG_ID(A5), ARG_ID(A6)); } )
|
||||
#define DEF7(N, A1, A2, A3, A4, A5, A6, A7) DEF7_F(N, A1, A2, A3, A4, A5, A6, A7, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3), ARG_ID(A4), ARG_ID(A5), ARG_ID(A6), ARG_ID(A7)); } )
|
||||
#define DEF8(N, A1, A2, A3, A4, A5, A6, A7, A8) DEF8_F(N, A1, A2, A3, A4, A5, A6, A7, A8, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3), ARG_ID(A4), ARG_ID(A5), ARG_ID(A6), ARG_ID(A7), ARG_ID(A8)); } )
|
||||
#define DEF9(N, A1, A2, A3, A4, A5, A6, A7, A8, A9) DEF9_F(N, A1, A2, A3, A4, A5, A6, A7, A8, A9, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3), ARG_ID(A4), ARG_ID(A5), ARG_ID(A6), ARG_ID(A7), ARG_ID(A8), ARG_ID(A9)); } )
|
||||
#define DEF10(N, A1, A2, A3, A4, A5, A6, A7, A8, A9, A10) DEF10_F(N, A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3), ARG_ID(A4), ARG_ID(A5), ARG_ID(A6), ARG_ID(A7), ARG_ID(A8), ARG_ID(A9), ARG_ID(A10)); } )
|
||||
|
||||
// NOLINTEND(whitespace/line_length)
|
||||
|
||||
#endif // MUJOCO_PYTHON_EXPERIMENTAL_DEAR_IMGUI_DEAR_IMGUI_MACROS_H_
|
||||
@@ -0,0 +1,438 @@
|
||||
// Copyright 2026 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
|
||||
//
|
||||
// https://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.
|
||||
|
||||
#define NAMESPACE ImPlot
|
||||
#include "dear_imgui_macros.h"
|
||||
#include <implot.h>
|
||||
#include <pybind11/eval.h>
|
||||
#include <pybind11/pybind11.h>
|
||||
#include <pybind11/pytypes.h>
|
||||
#include <pybind11/stl.h>
|
||||
|
||||
// NOLINTBEGIN(whitespace/line_length)
|
||||
|
||||
namespace py = pybind11;
|
||||
using ImString = const char*;
|
||||
static constexpr const ImVec2 ImVec2_Zero = ImVec2(0.0f, 0.0f);
|
||||
static constexpr const ImVec2 ImVec2_One = ImVec2(1.0f, 1.0f);
|
||||
static constexpr const ImVec2 ImVec2_NegOne_Zero = ImVec2(-1.0f, 0.0f);
|
||||
static constexpr const ImVec4 ImVec4_Zero = ImVec4(0.0f, 0.0f, 0.0f, 0.0f);
|
||||
static constexpr const ImVec4 ImVec4_One = ImVec4(1.0f, 1.0f, 1.0f, 1.0f);
|
||||
static constexpr const ImPlotRect ImPlotRect_Default{};
|
||||
static constexpr const ImPlotRange ImPlotRange_Default{};
|
||||
|
||||
PYBIND11_MODULE(implot, m) {
|
||||
// Import dear_imgui to make types like ImVec2 available.
|
||||
py::module_::import("mujoco.experimental.dear_imgui.dear_imgui");
|
||||
|
||||
// Types.
|
||||
py::class_<ImPlotPoint>(m, "Point")
|
||||
.def(py::init<>())
|
||||
.def(py::init<double, double>(), py::arg("_x"), py::arg("_y"))
|
||||
.def_readwrite("x", &ImPlotPoint::x)
|
||||
.def_readwrite("y", &ImPlotPoint::y);
|
||||
|
||||
py::class_<ImPlotRange>(m, "Range")
|
||||
.def(py::init<>())
|
||||
.def(py::init<double, double>(), py::arg("_min"), py::arg("_max"))
|
||||
.def_readwrite("min", &ImPlotRange::Min)
|
||||
.def_readwrite("max", &ImPlotRange::Max);
|
||||
|
||||
// py::class_<ImPlotRect>(m, "Rect")
|
||||
// .def(py::init<>())
|
||||
// .def(py::init<ImPlotRange, ImPlotRange>(), py::arg("_x"), py::arg("_y"))
|
||||
// .def_readwrite("x", &ImPlotRect::X)
|
||||
// .def_readwrite("y", &ImPlotRect::Y);
|
||||
|
||||
// Enumerations.
|
||||
|
||||
py::enum_<ImAxis_>(m, "Axis")
|
||||
.value("X1", ImAxis_X1)
|
||||
.value("X2", ImAxis_X2)
|
||||
.value("X3", ImAxis_X3)
|
||||
.value("Y1", ImAxis_Y1)
|
||||
.value("Y2", ImAxis_Y2)
|
||||
.value("Y3", ImAxis_Y3);
|
||||
|
||||
py::enum_<ImPlotFlags_>(m, "Flags")
|
||||
.value("None", ImPlotFlags_None)
|
||||
.value("NoTitle", ImPlotFlags_NoTitle)
|
||||
.value("NoLegend", ImPlotFlags_NoLegend)
|
||||
.value("NoMouseText", ImPlotFlags_NoMouseText)
|
||||
.value("NoInputs", ImPlotFlags_NoInputs)
|
||||
.value("NoMenus", ImPlotFlags_NoMenus)
|
||||
.value("NoBoxSelect", ImPlotFlags_NoBoxSelect)
|
||||
.value("NoFrame", ImPlotFlags_NoFrame)
|
||||
.value("Equal", ImPlotFlags_Equal)
|
||||
.value("Crosshairs", ImPlotFlags_Crosshairs)
|
||||
.value("CanvasOnly", ImPlotFlags_CanvasOnly);
|
||||
|
||||
py::enum_<ImPlotAxisFlags_>(m, "AxisFlags")
|
||||
.value("None", ImPlotAxisFlags_None)
|
||||
.value("NoLabel", ImPlotAxisFlags_NoLabel)
|
||||
.value("NoGridLines", ImPlotAxisFlags_NoGridLines)
|
||||
.value("NoTickMarks", ImPlotAxisFlags_NoTickMarks)
|
||||
.value("NoTickLabels", ImPlotAxisFlags_NoTickLabels)
|
||||
.value("NoInitialFit", ImPlotAxisFlags_NoInitialFit)
|
||||
.value("NoMenus", ImPlotAxisFlags_NoMenus)
|
||||
.value("NoSideSwitch", ImPlotAxisFlags_NoSideSwitch)
|
||||
.value("NoHighlight", ImPlotAxisFlags_NoHighlight)
|
||||
.value("Opposite", ImPlotAxisFlags_Opposite)
|
||||
.value("Foreground", ImPlotAxisFlags_Foreground)
|
||||
.value("Invert", ImPlotAxisFlags_Invert)
|
||||
.value("AutoFit", ImPlotAxisFlags_AutoFit)
|
||||
.value("RangeFit", ImPlotAxisFlags_RangeFit)
|
||||
.value("PanStretch", ImPlotAxisFlags_PanStretch)
|
||||
.value("LockMin", ImPlotAxisFlags_LockMin)
|
||||
.value("LockMax", ImPlotAxisFlags_LockMax)
|
||||
.value("Lock", ImPlotAxisFlags_Lock)
|
||||
.value("NoDecorations", ImPlotAxisFlags_NoDecorations)
|
||||
.value("AuxDefault", ImPlotAxisFlags_AuxDefault);
|
||||
|
||||
py::enum_<ImPlotSubplotFlags_>(m, "SubplotFlags")
|
||||
.value("None", ImPlotSubplotFlags_None)
|
||||
.value("NoTitle", ImPlotSubplotFlags_NoTitle)
|
||||
.value("NoLegend", ImPlotSubplotFlags_NoLegend)
|
||||
.value("NoMenus", ImPlotSubplotFlags_NoMenus)
|
||||
.value("NoResize", ImPlotSubplotFlags_NoResize)
|
||||
.value("NoAlign", ImPlotSubplotFlags_NoAlign)
|
||||
.value("ShareItems", ImPlotSubplotFlags_ShareItems)
|
||||
.value("LinkRows", ImPlotSubplotFlags_LinkRows)
|
||||
.value("LinkCols", ImPlotSubplotFlags_LinkCols)
|
||||
.value("LinkAllX", ImPlotSubplotFlags_LinkAllX)
|
||||
.value("LinkAllY", ImPlotSubplotFlags_LinkAllY)
|
||||
.value("ColMajor", ImPlotSubplotFlags_ColMajor);
|
||||
|
||||
py::enum_<ImPlotLegendFlags_>(m, "LegendFlags")
|
||||
.value("None", ImPlotLegendFlags_None)
|
||||
.value("NoButtons", ImPlotLegendFlags_NoButtons)
|
||||
.value("NoHighlightItem", ImPlotLegendFlags_NoHighlightItem)
|
||||
.value("NoHighlightAxis", ImPlotLegendFlags_NoHighlightAxis)
|
||||
.value("NoMenus", ImPlotLegendFlags_NoMenus)
|
||||
.value("Outside", ImPlotLegendFlags_Outside)
|
||||
.value("Horizontal", ImPlotLegendFlags_Horizontal)
|
||||
.value("Sort", ImPlotLegendFlags_Sort)
|
||||
.value("Reverse", ImPlotLegendFlags_Reverse);
|
||||
|
||||
py::enum_<ImPlotMouseTextFlags_>(m, "MouseTextFlags")
|
||||
.value("None", ImPlotMouseTextFlags_None)
|
||||
.value("NoAuxAxes", ImPlotMouseTextFlags_NoAuxAxes)
|
||||
.value("NoFormat", ImPlotMouseTextFlags_NoFormat)
|
||||
.value("ShowAlways", ImPlotMouseTextFlags_ShowAlways);
|
||||
|
||||
py::enum_<ImPlotDragToolFlags_>(m, "DragToolFlags")
|
||||
.value("None", ImPlotDragToolFlags_None)
|
||||
.value("NoCursors", ImPlotDragToolFlags_NoCursors)
|
||||
.value("NoFit", ImPlotDragToolFlags_NoFit)
|
||||
.value("NoInputs", ImPlotDragToolFlags_NoInputs)
|
||||
.value("Delayed", ImPlotDragToolFlags_Delayed);
|
||||
|
||||
py::enum_<ImPlotColormapScaleFlags_>(m, "ColormapScaleFlags")
|
||||
.value("None", ImPlotColormapScaleFlags_None)
|
||||
.value("NoLabel", ImPlotColormapScaleFlags_NoLabel)
|
||||
.value("Opposite", ImPlotColormapScaleFlags_Opposite)
|
||||
.value("Invert", ImPlotColormapScaleFlags_Invert);
|
||||
|
||||
py::enum_<ImPlotItemFlags_>(m, "ItemFlags")
|
||||
.value("None", ImPlotItemFlags_None)
|
||||
.value("NoLegend", ImPlotItemFlags_NoLegend)
|
||||
.value("NoFit", ImPlotItemFlags_NoFit);
|
||||
|
||||
py::enum_<ImPlotLineFlags_>(m, "LineFlags")
|
||||
.value("None", ImPlotLineFlags_None)
|
||||
.value("Segments", ImPlotLineFlags_Segments)
|
||||
.value("Loop", ImPlotLineFlags_Loop)
|
||||
.value("SkipNaN", ImPlotLineFlags_SkipNaN)
|
||||
.value("NoClip", ImPlotLineFlags_NoClip)
|
||||
.value("Shaded", ImPlotLineFlags_Shaded);
|
||||
|
||||
py::enum_<ImPlotScatterFlags_>(m, "ScatterFlags")
|
||||
.value("None", ImPlotScatterFlags_None)
|
||||
.value("NoClip", ImPlotScatterFlags_NoClip);
|
||||
|
||||
py::enum_<ImPlotStairsFlags_>(m, "StairsFlags")
|
||||
.value("None", ImPlotStairsFlags_None)
|
||||
.value("PreStep", ImPlotStairsFlags_PreStep)
|
||||
.value("Shaded", ImPlotStairsFlags_Shaded);
|
||||
|
||||
py::enum_<ImPlotShadedFlags_>(m, "ShadedFlags")
|
||||
.value("None", ImPlotShadedFlags_None);
|
||||
|
||||
py::enum_<ImPlotBarsFlags_>(m, "BarsFlags")
|
||||
.value("None", ImPlotBarsFlags_None)
|
||||
.value("Horizontal", ImPlotBarsFlags_Horizontal);
|
||||
|
||||
py::enum_<ImPlotBarGroupsFlags_>(m, "BarGroupsFlags")
|
||||
.value("None", ImPlotBarGroupsFlags_None)
|
||||
.value("Horizontal", ImPlotBarGroupsFlags_Horizontal)
|
||||
.value("Stacked", ImPlotBarGroupsFlags_Stacked);
|
||||
|
||||
py::enum_<ImPlotErrorBarsFlags_>(m, "ErrorBarsFlags")
|
||||
.value("None", ImPlotErrorBarsFlags_None)
|
||||
.value("Horizontal", ImPlotErrorBarsFlags_Horizontal);
|
||||
|
||||
py::enum_<ImPlotStemsFlags_>(m, "StemsFlags")
|
||||
.value("None", ImPlotStemsFlags_None)
|
||||
.value("Horizontal", ImPlotStemsFlags_Horizontal);
|
||||
|
||||
py::enum_<ImPlotInfLinesFlags_>(m, "InfLinesFlags")
|
||||
.value("None", ImPlotInfLinesFlags_None)
|
||||
.value("Horizontal", ImPlotInfLinesFlags_Horizontal);
|
||||
|
||||
py::enum_<ImPlotPieChartFlags_>(m, "PieChartFlags")
|
||||
.value("None", ImPlotPieChartFlags_None)
|
||||
.value("Normalize", ImPlotPieChartFlags_Normalize)
|
||||
.value("IgnoreHidden", ImPlotPieChartFlags_IgnoreHidden)
|
||||
.value("Exploding", ImPlotPieChartFlags_Exploding);
|
||||
|
||||
py::enum_<ImPlotHeatmapFlags_>(m, "HeatmapFlags")
|
||||
.value("None", ImPlotHeatmapFlags_None)
|
||||
.value("ColMajor", ImPlotHeatmapFlags_ColMajor);
|
||||
|
||||
py::enum_<ImPlotHistogramFlags_>(m, "HistogramFlags")
|
||||
.value("None", ImPlotHistogramFlags_None)
|
||||
.value("Horizontal", ImPlotHistogramFlags_Horizontal)
|
||||
.value("Cumulative", ImPlotHistogramFlags_Cumulative)
|
||||
.value("Density", ImPlotHistogramFlags_Density)
|
||||
.value("NoOutliers", ImPlotHistogramFlags_NoOutliers)
|
||||
.value("ColMajor", ImPlotHistogramFlags_ColMajor);
|
||||
|
||||
py::enum_<ImPlotDigitalFlags_>(m, "DigitalFlags")
|
||||
.value("ImPlotNone", ImPlotDigitalFlags_None);
|
||||
|
||||
py::enum_<ImPlotImageFlags_>(m, "ImageFlags")
|
||||
.value("None", ImPlotImageFlags_None);
|
||||
|
||||
py::enum_<ImPlotTextFlags_>(m, "TextFlags")
|
||||
.value("None", ImPlotTextFlags_None)
|
||||
.value("Vertical", ImPlotTextFlags_Vertical);
|
||||
|
||||
py::enum_<ImPlotDummyFlags_>(m, "DummyFlags")
|
||||
.value("None", ImPlotDummyFlags_None);
|
||||
|
||||
py::enum_<ImPlotCond_>(m, "Cond")
|
||||
.value("None", ImPlotCond_None)
|
||||
.value("Always", ImPlotCond_Always)
|
||||
.value("Once", ImPlotCond_Once);
|
||||
|
||||
py::enum_<ImPlotCol_>(m, "Col")
|
||||
.value("Line", ImPlotCol_Line)
|
||||
.value("Fill", ImPlotCol_Fill)
|
||||
.value("MarkerOutline", ImPlotCol_MarkerOutline)
|
||||
.value("MarkerFill", ImPlotCol_MarkerFill)
|
||||
.value("ErrorBar", ImPlotCol_ErrorBar)
|
||||
.value("FrameBg", ImPlotCol_FrameBg)
|
||||
.value("PlotBg", ImPlotCol_PlotBg)
|
||||
.value("PlotBorder", ImPlotCol_PlotBorder)
|
||||
.value("LegendBg", ImPlotCol_LegendBg)
|
||||
.value("LegendBorder", ImPlotCol_LegendBorder)
|
||||
.value("LegendText", ImPlotCol_LegendText)
|
||||
.value("TitleText", ImPlotCol_TitleText)
|
||||
.value("InlayText", ImPlotCol_InlayText)
|
||||
.value("AxisText", ImPlotCol_AxisText)
|
||||
.value("AxisGrid", ImPlotCol_AxisGrid)
|
||||
.value("AxisTick", ImPlotCol_AxisTick)
|
||||
.value("AxisBg", ImPlotCol_AxisBg)
|
||||
.value("AxisBgHovered", ImPlotCol_AxisBgHovered)
|
||||
.value("AxisBgActive", ImPlotCol_AxisBgActive)
|
||||
.value("Selection", ImPlotCol_Selection)
|
||||
.value("Crosshairs", ImPlotCol_Crosshairs);
|
||||
|
||||
py::enum_<ImPlotStyleVar_>(m, "StyleVar")
|
||||
.value("LineWeight", ImPlotStyleVar_LineWeight)
|
||||
.value("Marker", ImPlotStyleVar_Marker)
|
||||
.value("MarkerSize", ImPlotStyleVar_MarkerSize)
|
||||
.value("MarkerWeight", ImPlotStyleVar_MarkerWeight)
|
||||
.value("FillAlpha", ImPlotStyleVar_FillAlpha)
|
||||
.value("ErrorBarSize", ImPlotStyleVar_ErrorBarSize)
|
||||
.value("ErrorBarWeight", ImPlotStyleVar_ErrorBarWeight)
|
||||
.value("DigitalBitHeight", ImPlotStyleVar_DigitalBitHeight)
|
||||
.value("DigitalBitGap", ImPlotStyleVar_DigitalBitGap)
|
||||
.value("PlotBorderSize", ImPlotStyleVar_PlotBorderSize)
|
||||
.value("MinorAlpha", ImPlotStyleVar_MinorAlpha)
|
||||
.value("MajorTickLen", ImPlotStyleVar_MajorTickLen)
|
||||
.value("MinorTickLen", ImPlotStyleVar_MinorTickLen)
|
||||
.value("MajorTickSize", ImPlotStyleVar_MajorTickSize)
|
||||
.value("MinorTickSize", ImPlotStyleVar_MinorTickSize)
|
||||
.value("MajorGridSize", ImPlotStyleVar_MajorGridSize)
|
||||
.value("MinorGridSize", ImPlotStyleVar_MinorGridSize)
|
||||
.value("PlotPadding", ImPlotStyleVar_PlotPadding)
|
||||
.value("LabelPadding", ImPlotStyleVar_LabelPadding)
|
||||
.value("LegendPadding", ImPlotStyleVar_LegendPadding)
|
||||
.value("LegendInnerPadding", ImPlotStyleVar_LegendInnerPadding)
|
||||
.value("LegendSpacing", ImPlotStyleVar_LegendSpacing)
|
||||
.value("MousePosPadding", ImPlotStyleVar_MousePosPadding)
|
||||
.value("AnnotationPadding", ImPlotStyleVar_AnnotationPadding)
|
||||
.value("FitPadding", ImPlotStyleVar_FitPadding)
|
||||
.value("PlotDefaultSize", ImPlotStyleVar_PlotDefaultSize)
|
||||
.value("PlotMinSize", ImPlotStyleVar_PlotMinSize);
|
||||
|
||||
py::enum_<ImPlotScale_>(m, "Scale")
|
||||
.value("ImPlotScale_Linear", ImPlotScale_Linear)
|
||||
.value("ImPlotScale_Time", ImPlotScale_Time)
|
||||
.value("ImPlotScale_Log10", ImPlotScale_Log10)
|
||||
.value("ImPlotScale_SymLog", ImPlotScale_SymLog);
|
||||
|
||||
py::enum_<ImPlotMarker_>(m, "Marker")
|
||||
.value("None", ImPlotMarker_None)
|
||||
.value("Circle", ImPlotMarker_Circle)
|
||||
.value("Square", ImPlotMarker_Square)
|
||||
.value("Diamond", ImPlotMarker_Diamond)
|
||||
.value("Up", ImPlotMarker_Up)
|
||||
.value("Down", ImPlotMarker_Down)
|
||||
.value("Left", ImPlotMarker_Left)
|
||||
.value("Right", ImPlotMarker_Right)
|
||||
.value("Cross", ImPlotMarker_Cross)
|
||||
.value("Plus", ImPlotMarker_Plus)
|
||||
.value("Asterisk", ImPlotMarker_Asterisk);
|
||||
|
||||
py::enum_<ImPlotLocation_>(m, "Location")
|
||||
.value("Center", ImPlotLocation_Center)
|
||||
.value("North", ImPlotLocation_North)
|
||||
.value("South", ImPlotLocation_South)
|
||||
.value("West", ImPlotLocation_West)
|
||||
.value("East", ImPlotLocation_East)
|
||||
.value("NorthWest", ImPlotLocation_NorthWest)
|
||||
.value("NorthEast", ImPlotLocation_NorthEast)
|
||||
.value("SouthWest", ImPlotLocation_SouthWest)
|
||||
.value("SouthEast", ImPlotLocation_SouthEast);
|
||||
|
||||
// Functions.
|
||||
|
||||
DEF3(BeginPlot, (ImString, title_id, ), (const ImVec2&, size, = ImVec2_NegOne_Zero), (ImPlotFlags, flags, = 0));
|
||||
DEF0(EndPlot);
|
||||
DEF7(BeginSubplots, (ImString, title_id, ), (int, rows, ), (int, cols, ), (const ImVec2&, size, ), (ImPlotSubplotFlags, flags, = 0), (float*, row_ratios, = nullptr), (float*, col_ratios, = nullptr));
|
||||
DEF0(EndSubplots);
|
||||
DEF3(SetupAxis, (ImAxis, axis, ), (ImString, label, = nullptr), (ImPlotAxisFlags, flags, = 0));
|
||||
DEF4(SetupAxisLimits, (ImAxis, axis, ), (double, v_min, ), (double, v_max, ), (ImPlotCond, cond, = ImPlotCond_Once));
|
||||
DEF4_F(SetupAxisTicks, (ImAxis, axis, ), (std::vector<double>, values, ), (std::vector<std::string>, labels, ), (bool, keep_default, = false), {
|
||||
std::vector<const char*> c_labels;
|
||||
c_labels.reserve(labels.size());
|
||||
for (const auto& l : labels) {
|
||||
c_labels.push_back(l.c_str());
|
||||
}
|
||||
ImPlot::SetupAxisTicks(axis, values.data(), values.size(), c_labels.empty() ? nullptr : c_labels.data(), keep_default);
|
||||
});
|
||||
DEF3(SetupAxisLinks, (ImAxis, axis, ), (double*, link_min, ), (double*, link_max, ));
|
||||
DEF2(SetupAxisFormat, (ImAxis, axis, ), (ImString, fmt, ));
|
||||
DEF2(SetupAxisScale, (ImAxis, axis, ), (ImPlotScale, scale, ));
|
||||
DEF3(SetupAxisLimitsConstraints, (ImAxis, axis, ), (double, v_min, ), (double, v_max, ));
|
||||
DEF3(SetupAxisZoomConstraints, (ImAxis, axis, ), (double, z_min, ), (double, z_max, ));
|
||||
DEF4(SetupAxes, (ImString, x_label, ), (ImString, y_label, ), (ImPlotAxisFlags, x_flags, = 0), (ImPlotAxisFlags, y_flags, = 0));
|
||||
DEF5(SetupAxesLimits, (double, x_min, ), (double, x_max, ), (double, y_min, ), (double, y_max, ), (ImPlotCond, cond, = ImPlotCond_Once));
|
||||
DEF2(SetupLegend, (ImPlotLocation, location, ), (ImPlotLegendFlags, flags, = 0));
|
||||
DEF2(SetupMouseText, (ImPlotLocation, location, ), (ImPlotMouseTextFlags, flags, = 0));
|
||||
DEF0(SetupFinish);
|
||||
DEF4(SetNextAxisLimits, (ImAxis, axis, ), (double, v_min, ), (double, v_max, ), (ImPlotCond, cond, = ImPlotCond_Once));
|
||||
DEF3(SetNextAxisLinks, (ImAxis, axis, ), (double*, link_min, ), (double*, link_max, ));
|
||||
DEF1(SetNextAxisToFit, (ImAxis, axis, ));
|
||||
DEF5(SetNextAxesLimits, (double, x_min, ), (double, x_max, ), (double, y_min, ), (double, y_max, ), (ImPlotCond, cond, = ImPlotCond_Once));
|
||||
DEF0_F(SetNextAxesToFit, {
|
||||
return ImPlot::SetNextAxesToFit();
|
||||
});
|
||||
DEF6_F(PlotLine, (ImString, label_id, ), (std::vector<double>, xs, ), (std::vector<double>, ys, ), (ImPlotLineFlags, flags, = 0), (int, offset, = 0), (int, stride, = sizeof(double)), {
|
||||
return ImPlot::PlotLine(label_id, xs.data(), ys.data(), xs.size(), flags, offset, stride);
|
||||
});
|
||||
DEF6_F(PlotScatter, (ImString, label_id, ), (std::vector<double>, xs, ), (std::vector<double>, ys, ), (ImPlotScatterFlags, flags, = 0), (int, offset, = 0), (int, stride, = sizeof(double)), {
|
||||
return ImPlot::PlotScatter(label_id, xs.data(), ys.data(), xs.size(), flags, offset, stride);
|
||||
});
|
||||
DEF6_F(PlotStairs, (ImString, label_id, ), (std::vector<double>, xs, ), (std::vector<double>, ys, ), (ImPlotStairsFlags, flags, = 0), (int, offset, = 0), (int, stride, = sizeof(double)), {
|
||||
return ImPlot::PlotStairs(label_id, xs.data(), ys.data(), xs.size(), flags, offset, stride);
|
||||
});
|
||||
DEF7_F(PlotShaded, (ImString, label_id, ), (std::vector<double>, xs, ), (std::vector<double>, ys, ), (double, yref, = 0), (ImPlotShadedFlags, flags, = 0), (int, offset, = 0), (int, stride, = sizeof(double)), {
|
||||
return ImPlot::PlotShaded(label_id, xs.data(), ys.data(), xs.size(), yref, flags, offset, stride);
|
||||
});
|
||||
DEF7_F(PlotShaded, (ImString, label_id, ), (std::vector<double>, xs, ), (std::vector<double>, ys1, ), (std::vector<double>, ys2, ), (ImPlotShadedFlags, flags, = 0), (int, offset, = 0), (int, stride, = sizeof(double)), {
|
||||
return ImPlot::PlotShaded(label_id, xs.data(), ys1.data(), ys2.data(), xs.size(), flags, offset, stride);
|
||||
});
|
||||
DEF7_F(PlotBars, (ImString, label_id, ), (std::vector<double>, xs, ), (std::vector<double>, ys, ), (double, bar_size, ), (ImPlotBarsFlags, flags, = 0), (int, offset, = 0), (int, stride, = sizeof(double)), {
|
||||
return ImPlot::PlotBars(label_id, xs.data(), ys.data(), xs.size(), bar_size, flags, offset, stride);
|
||||
});
|
||||
DEF7_F(PlotErrorBars, (ImString, label_id, ), (std::vector<double>, xs, ), (std::vector<double>, ys, ), (std::vector<double>, err, ), (ImPlotErrorBarsFlags, flags, = 0), (int, offset, = 0), (int, stride, = sizeof(double)), {
|
||||
return ImPlot::PlotErrorBars(label_id, xs.data(), ys.data(), err.data(), xs.size(), flags, offset, stride);
|
||||
});
|
||||
DEF8_F(PlotErrorBars, (ImString, label_id, ), (std::vector<double>, xs, ), (std::vector<double>, ys, ), (std::vector<double>, neg, ), (std::vector<double>, pos, ), (ImPlotErrorBarsFlags, flags, = 0), (int, offset, = 0), (int, stride, = sizeof(double)), {
|
||||
return ImPlot::PlotErrorBars(label_id, xs.data(), ys.data(), neg.data(), pos.data(), xs.size(), flags, offset, stride);
|
||||
});
|
||||
DEF7_F(PlotStems, (ImString, label_id, ), (std::vector<double>, xs, ), (std::vector<double>, ys, ), (double, ref, = 0), (ImPlotStemsFlags, flags, = 0), (int, offset, = 0), (int, stride, = sizeof(double)), {
|
||||
return ImPlot::PlotStems(label_id, xs.data(), ys.data(), xs.size(), ref, flags, offset, stride);
|
||||
});
|
||||
DEF6_F(PlotDigital, (ImString, label_id, ), (std::vector<double>, xs, ), (std::vector<double>, ys, ), (ImPlotDigitalFlags, flags, = 0), (int, offset, = 0), (int, stride, = sizeof(double)), {
|
||||
return ImPlot::PlotDigital(label_id, xs.data(), ys.data(), xs.size(), flags, offset, stride);
|
||||
});
|
||||
DEF8(PlotImage, (ImString, label_id, ), (ImTextureRef, tex_ref, ), (const ImPlotPoint&, bounds_min, ), (const ImPlotPoint&, bounds_max, ), (const ImVec2&, uv0, = ImVec2_Zero), (const ImVec2&, uv1, = ImVec2_One), (const ImVec4&, tint_col, = ImVec4_One), (ImPlotImageFlags, flags, = 0));
|
||||
DEF5(PlotText, (ImString, text, ), (double, x, ), (double, y, ), (const ImVec2&, pix_offset, = ImVec2_Zero), (ImPlotTextFlags, flags, = 0));
|
||||
DEF2(PlotDummy, (ImString, label_id, ), (ImPlotDummyFlags, flags, = 0));
|
||||
DEF9(DragPoint, (int, id, ), (double*, x, ), (double*, y, ), (const ImVec4&, col, ), (float, size, = 4), (ImPlotDragToolFlags, flags, = 0), (bool*, out_clicked, = nullptr), (bool*, out_hovered, = nullptr), (bool*, out_held, = nullptr));
|
||||
DEF8(DragLineX, (int, id, ), (double*, x, ), (const ImVec4&, col, ), (float, thickness, = 1), (ImPlotDragToolFlags, flags, = 0), (bool*, out_clicked, = nullptr), (bool*, out_hovered, = nullptr), (bool*, out_held, = nullptr));
|
||||
DEF8(DragLineY, (int, id, ), (double*, y, ), (const ImVec4&, col, ), (float, thickness, = 1), (ImPlotDragToolFlags, flags, = 0), (bool*, out_clicked, = nullptr), (bool*, out_hovered, = nullptr), (bool*, out_held, = nullptr));
|
||||
DEF10(DragRect, (int, id, ), (double*, x1, ), (double*, y1, ), (double*, x2, ), (double*, y2, ), (const ImVec4&, col, ), (ImPlotDragToolFlags, flags, = 0), (bool*, out_clicked, = nullptr), (bool*, out_hovered, = nullptr), (bool*, out_held, = nullptr));
|
||||
DEF6(Annotation, (double, x, ), (double, y, ), (const ImVec4&, col, ), (const ImVec2&, pix_offset, ), (bool, clamp, ), (bool, round, = false));
|
||||
DEF6_F(Annotation, (double, x, ), (double, y, ), (const ImVec4&, col, ), (const ImVec2&, pix_offset, ), (bool, clamp, ), (ImString, txt, ), {
|
||||
return ImPlot::Annotation(x, y, col, pix_offset, clamp, "%s", txt);
|
||||
});
|
||||
DEF3(TagX, (double, x, ), (const ImVec4&, col, ), (bool, round, = false));
|
||||
DEF3_F(TagX, (double, x, ), (const ImVec4&, col, ), (ImString, txt, ), {
|
||||
return ImPlot::TagX(x, col, "%s", txt);
|
||||
});
|
||||
DEF3(TagY, (double, y, ), (const ImVec4&, col, ), (bool, round, = false));
|
||||
DEF3_F(TagY, (double, y, ), (const ImVec4&, col, ), (ImString, txt, ), {
|
||||
return ImPlot::TagY(y, col, "%s", txt);
|
||||
});
|
||||
DEF1(SetAxis, (ImAxis, axis, ));
|
||||
DEF2(SetAxes, (ImAxis, x_axis, ), (ImAxis, y_axis, ));
|
||||
DEF3(PixelsToPlot, (const ImVec2&, pix, ), (ImAxis, x_axis, = IMPLOT_AUTO), (ImAxis, y_axis, = IMPLOT_AUTO));
|
||||
DEF4(PixelsToPlot, (float, x, ), (float, y, ), (ImAxis, x_axis, = IMPLOT_AUTO), (ImAxis, y_axis, = IMPLOT_AUTO));
|
||||
DEF3(PlotToPixels, (const ImPlotPoint&, plt, ), (ImAxis, x_axis, = IMPLOT_AUTO), (ImAxis, y_axis, = IMPLOT_AUTO));
|
||||
DEF4(PlotToPixels, (double, x, ), (double, y, ), (ImAxis, x_axis, = IMPLOT_AUTO), (ImAxis, y_axis, = IMPLOT_AUTO));
|
||||
DEF0(GetPlotPos);
|
||||
DEF0(GetPlotSize);
|
||||
DEF2(GetPlotMousePos, (ImAxis, x_axis, = IMPLOT_AUTO), (ImAxis, y_axis, = IMPLOT_AUTO));
|
||||
DEF2(GetPlotLimits, (ImAxis, x_axis, = IMPLOT_AUTO), (ImAxis, y_axis, = IMPLOT_AUTO));
|
||||
DEF0(IsPlotHovered);
|
||||
DEF1(IsAxisHovered, (ImAxis, axis, ));
|
||||
DEF0(IsSubplotsHovered);
|
||||
DEF0(IsPlotSelected);
|
||||
DEF2(GetPlotSelection, (ImAxis, x_axis, = IMPLOT_AUTO), (ImAxis, y_axis, = IMPLOT_AUTO));
|
||||
DEF0(CancelPlotSelection);
|
||||
DEF2(HideNextItem, (bool, hidden, = true), (ImPlotCond, cond, = ImPlotCond_Once));
|
||||
DEF2(BeginAlignedPlots, (ImString, group_id, ), (bool, vertical, = true));
|
||||
DEF0(EndAlignedPlots);
|
||||
DEF2(BeginLegendPopup, (ImString, label_id, ), (ImGuiMouseButton, mouse_button, = 1));
|
||||
DEF0(EndLegendPopup);
|
||||
DEF1(IsLegendEntryHovered, (ImString, label_id, ));
|
||||
DEF0(BeginDragDropTargetPlot);
|
||||
DEF1(BeginDragDropTargetAxis, (ImAxis, axis, ));
|
||||
DEF0(BeginDragDropTargetLegend);
|
||||
DEF0(EndDragDropTarget);
|
||||
DEF1(BeginDragDropSourcePlot, (ImGuiDragDropFlags, flags, = 0));
|
||||
DEF2(BeginDragDropSourceAxis, (ImAxis, axis, ), (ImGuiDragDropFlags, flags, = 0));
|
||||
DEF2(BeginDragDropSourceItem, (ImString, label_id, ), (ImGuiDragDropFlags, flags, = 0));
|
||||
DEF0(EndDragDropSource);
|
||||
DEF2(PushStyleColor, (ImPlotCol, idx, ), (ImU32, col, ));
|
||||
DEF2(PushStyleColor, (ImPlotCol, idx, ), (const ImVec4&, col, ));
|
||||
DEF1(PopStyleColor, (int, count, = 1));
|
||||
DEF2(PushStyleVar, (ImPlotStyleVar, idx, ), (float, val, ));
|
||||
DEF2(PushStyleVar, (ImPlotStyleVar, idx, ), (int, val, ));
|
||||
DEF2(PushStyleVar, (ImPlotStyleVar, idx, ), (const ImVec2&, val, ));
|
||||
DEF1(PopStyleVar, (int, count, = 1));
|
||||
DEF2(SetNextLineStyle, (const ImVec4&, col, = IMPLOT_AUTO_COL), (float, weight, = IMPLOT_AUTO));
|
||||
DEF2(SetNextFillStyle, (const ImVec4&, col, = IMPLOT_AUTO_COL), (float, alpha_mod, = IMPLOT_AUTO));
|
||||
DEF5(SetNextMarkerStyle, (ImPlotMarker, marker, = IMPLOT_AUTO), (float, size, = IMPLOT_AUTO), (const ImVec4&, fill, = IMPLOT_AUTO_COL), (float, weight, = IMPLOT_AUTO), (const ImVec4&, outline, = IMPLOT_AUTO_COL));
|
||||
DEF3(SetNextErrorBarStyle, (const ImVec4&, col, = IMPLOT_AUTO_COL), (float, size, = IMPLOT_AUTO), (float, weight, = IMPLOT_AUTO));
|
||||
DEF1(PushPlotClipRect, (float, expand, = 0));
|
||||
DEF0(PopPlotClipRect);
|
||||
}
|
||||
|
||||
// NOLINTEND(whitespace/line_length)
|
||||
@@ -0,0 +1,189 @@
|
||||
// Copyright 2026 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
|
||||
//
|
||||
// https://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 <cstddef>
|
||||
#include <cstdint>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <string_view>
|
||||
#include <vector>
|
||||
|
||||
#include <fstream>
|
||||
#include <implot.h>
|
||||
#include <mujoco/mujoco.h>
|
||||
#include "third_party/mujoco/src/experimental/platform/hal/graphics_mode.h"
|
||||
#include "third_party/mujoco/src/experimental/platform/hal/renderer.h"
|
||||
#include "third_party/mujoco/src/experimental/platform/hal/window.h"
|
||||
#include "structs.h"
|
||||
#include <pybind11/eval.h>
|
||||
#include <pybind11/pybind11.h>
|
||||
#include <pybind11/pytypes.h>
|
||||
#include <pybind11/stl.h>
|
||||
|
||||
static bool IsCuda() {
|
||||
#ifdef CUDA
|
||||
return true;
|
||||
#else
|
||||
return false;
|
||||
#endif
|
||||
}
|
||||
|
||||
static bool IsCrd() {
|
||||
const char* display = getenv("DISPLAY");
|
||||
return display ? strcmp(display, ":20") == 0 : false;
|
||||
}
|
||||
|
||||
static std::vector<std::byte> LoadAsset(std::string_view path) {
|
||||
std::string file_path = "assets/" +
|
||||
std::string(path.substr(path.find(':') + 1));
|
||||
|
||||
std::ifstream file(file_path, std::ios::binary | std::ios::ate);
|
||||
if (!file.is_open()) {
|
||||
return {};
|
||||
}
|
||||
auto file_size = file.tellg();
|
||||
file.seekg(0, std::ios::beg);
|
||||
std::vector<std::byte> buffer(file_size);
|
||||
if (!file.read(reinterpret_cast<char*>(buffer.data()), file_size)) {
|
||||
return {};
|
||||
}
|
||||
return buffer;
|
||||
}
|
||||
|
||||
// Holds loaded resource data for the MuJoCo resource provider.
|
||||
struct ResourceData {
|
||||
std::vector<std::byte> bytes;
|
||||
};
|
||||
|
||||
class Viewer {
|
||||
public:
|
||||
Viewer(const std::string& title, int width, int height,
|
||||
std::string graphics_mode_str) {
|
||||
// Register resource providers for font and filament assets.
|
||||
mjpResourceProvider resource_provider;
|
||||
mjp_defaultResourceProvider(&resource_provider);
|
||||
|
||||
resource_provider.open = [](mjResource* resource) {
|
||||
auto* data = new ResourceData();
|
||||
data->bytes = LoadAsset(resource->name);
|
||||
resource->data = data;
|
||||
return static_cast<int>(data->bytes.size());
|
||||
};
|
||||
resource_provider.read = [](mjResource* resource, const void** buffer) {
|
||||
auto* data = static_cast<ResourceData*>(resource->data);
|
||||
*buffer = data->bytes.data();
|
||||
return static_cast<int>(data->bytes.size());
|
||||
};
|
||||
resource_provider.close = [](mjResource* resource) {
|
||||
delete static_cast<ResourceData*>(resource->data);
|
||||
resource->data = nullptr;
|
||||
};
|
||||
resource_provider.prefix = "font";
|
||||
mjp_registerResourceProvider(&resource_provider);
|
||||
resource_provider.prefix = "filament";
|
||||
mjp_registerResourceProvider(&resource_provider);
|
||||
|
||||
mujoco::platform::Window::Config config;
|
||||
using GraphicsMode = mujoco::platform::GraphicsMode;
|
||||
config.gfx_mode = mujoco::platform::GraphicsModeFromString(
|
||||
graphics_mode_str, GraphicsMode::FilamentOpenGl);
|
||||
window_ = std::make_unique<mujoco::platform::Window>("PyStudio " + title,
|
||||
width, height, config);
|
||||
ImPlot::CreateContext();
|
||||
|
||||
renderer_ = std::make_unique<mujoco::platform::Renderer>(
|
||||
window_->GetNativeWindowHandle(), config.gfx_mode);
|
||||
}
|
||||
|
||||
void InitRenderer(const mujoco::python::MjModelWrapper& model) {
|
||||
renderer_->Init(model.get());
|
||||
}
|
||||
|
||||
bool NewFrame() {
|
||||
const mujoco::platform::Window::Status status = window_->NewFrame();
|
||||
return status == mujoco::platform::Window::Status::kRunning;
|
||||
}
|
||||
|
||||
intptr_t UploadImage(intptr_t tex_id, const std::string img, int width,
|
||||
int height, int bpp) {
|
||||
return renderer_->UploadImage(tex_id, (const std::byte*)img.data(), width,
|
||||
height, bpp);
|
||||
}
|
||||
|
||||
int RenderToTexture(const mujoco::python::MjModelWrapper& model,
|
||||
mujoco::python::MjDataWrapper& data,
|
||||
mujoco::python::MjvCameraWrapper& cam, int width,
|
||||
int height, int tex_id) {
|
||||
const int bytes_per_pixel = 3;
|
||||
std::vector<std::byte> bytes(width * height * bytes_per_pixel);
|
||||
renderer_->RenderToTexture(model.get(), data.get(), cam.get(), width,
|
||||
height, bytes.data());
|
||||
return renderer_->UploadImage(tex_id, bytes.data(), width, height,
|
||||
bytes_per_pixel);
|
||||
}
|
||||
|
||||
std::string GetDropFile() {
|
||||
return window_->GetDropFile();
|
||||
}
|
||||
|
||||
void Present(const mujoco::python::MjModelWrapper& model,
|
||||
mujoco::python::MjDataWrapper& data,
|
||||
mujoco::python::MjvPerturbWrapper& perturb,
|
||||
mujoco::python::MjvCameraWrapper& camera,
|
||||
mujoco::python::MjvOptionWrapper& vis_options,
|
||||
const std::vector<uint8_t>& render_flags) {
|
||||
const float width = window_->GetWidth();
|
||||
const float height = window_->GetHeight();
|
||||
const float scale = window_->GetScale();
|
||||
|
||||
if (mujoco::platform::IsHeadless(window_->GetGraphicsMode())) {
|
||||
pixels_.resize(width * height * 3);
|
||||
} else {
|
||||
pixels_.clear();
|
||||
}
|
||||
|
||||
// Update render flags before rendering.
|
||||
mjtByte* flags = renderer_->GetRenderFlags();
|
||||
for (size_t i = 0; i < mjNRNDFLAG && i < render_flags.size(); ++i) {
|
||||
flags[i] = render_flags[i];
|
||||
}
|
||||
|
||||
renderer_->Render(model.get(), data.get(), perturb.get(), camera.get(),
|
||||
vis_options.get(), width * scale, height * scale,
|
||||
pixels_);
|
||||
|
||||
window_->EndFrame();
|
||||
window_->Present(pixels_);
|
||||
}
|
||||
|
||||
private:
|
||||
std::unique_ptr<mujoco::platform::Window> window_;
|
||||
std::unique_ptr<mujoco::platform::Renderer> renderer_;
|
||||
std::vector<std::byte> pixels_;
|
||||
};
|
||||
|
||||
PYBIND11_MODULE(native_viewer_cc, m) {
|
||||
pybind11::class_<Viewer>(m, "Viewer")
|
||||
.def(pybind11::init<const std::string&, int, int, const std::string&>())
|
||||
.def("InitRenderer", &Viewer::InitRenderer)
|
||||
.def("NewFrame", &Viewer::NewFrame)
|
||||
.def("Present", &Viewer::Present)
|
||||
.def("UploadImage", &Viewer::UploadImage)
|
||||
.def("RenderToTexture", &Viewer::RenderToTexture)
|
||||
.def("GetDropFile", &Viewer::GetDropFile);
|
||||
m.def("IsCrd", &IsCrd);
|
||||
m.def("IsCuda", &IsCuda);
|
||||
}
|
||||
@@ -0,0 +1,171 @@
|
||||
# Copyright 2026 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
|
||||
#
|
||||
# https://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.
|
||||
"""Simulation-agnostic native viewer for MuJoCo models.
|
||||
|
||||
This class is simulation-agnostic and as such it does not own the model or data.
|
||||
|
||||
See the documentation for studio_app.py for more details on the architecture
|
||||
separating the viewer and simulation. See the sample/ folder for examples of how
|
||||
to use these classes.
|
||||
"""
|
||||
|
||||
import mujoco
|
||||
|
||||
from mujoco.experimental.studio import native_viewer_cc as _viewer
|
||||
from mujoco.experimental.studio import ux
|
||||
|
||||
|
||||
class NativeViewer:
|
||||
"""Simulation-agnostic native viewer for MuJoCo models."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: mujoco.MjModel,
|
||||
camera: mujoco.MjvCamera | None = None,
|
||||
vis_options: mujoco.MjvOption | None = None,
|
||||
perturb: mujoco.MjvPerturb | None = None,
|
||||
render_flags: ux.RenderFlags | None = None,
|
||||
title: str = '',
|
||||
width: int = 1200,
|
||||
height: int = 800,
|
||||
gfx: str = '',
|
||||
) -> None:
|
||||
"""Initializes the NativeViewer.
|
||||
|
||||
The viewer creates and modifies its own camera, perturbation, and
|
||||
visualization option objects unless they are provided.
|
||||
|
||||
Args:
|
||||
model: The MuJoCo model, used to initialize the renderer.
|
||||
camera: Camera parameters. Internal object is created if None.
|
||||
vis_options: Visualization options. Internal object is created if None.
|
||||
perturb: Perturbation parameters. Internal object is created if None.
|
||||
render_flags: Render flags. Internal object is created if None.
|
||||
title: Title of the viewer window.
|
||||
width: Initial width of the viewer window.
|
||||
height: Initial height of the viewer window.
|
||||
gfx: Graphics mode.
|
||||
"""
|
||||
self.camera = camera or mujoco.MjvCamera()
|
||||
self.perturb = perturb or mujoco.MjvPerturb()
|
||||
self.vis_options = vis_options or mujoco.MjvOption()
|
||||
self._viewer = _viewer.Viewer(title, width, height, gfx)
|
||||
self._viewer.InitRenderer(model)
|
||||
# This class does not own the model but we need to know if the model being
|
||||
# rendered has changed, so we store the unique python object id here so we
|
||||
# can use it to detect model changes.
|
||||
self._renderer_model_id = id(model)
|
||||
self._is_running = True
|
||||
if render_flags is not None:
|
||||
self.render_flags = render_flags
|
||||
else:
|
||||
self.render_flags = ux.RenderFlags()
|
||||
# Initted to match mujoco/src/engine/engine_vis_init.c
|
||||
self.render_flags.flags = [1, 0, 1, 0, 1, 0, 1, 0, 0, 0, 1]
|
||||
|
||||
def _sync_renderer(self, model: mujoco.MjModel) -> None:
|
||||
"""Re-initializes the renderer if the model object has changed."""
|
||||
if id(model) != self._renderer_model_id:
|
||||
self._viewer.InitRenderer(model)
|
||||
self._renderer_model_id = id(model)
|
||||
|
||||
def is_running(self) -> bool:
|
||||
"""Poll for a new frame; returns ``False`` when the window is closed."""
|
||||
if not self._is_running:
|
||||
return False
|
||||
self._is_running = self._viewer.NewFrame()
|
||||
return self._is_running
|
||||
|
||||
def sync(
|
||||
self,
|
||||
model: mujoco.MjModel,
|
||||
data: mujoco.MjData,
|
||||
) -> None:
|
||||
"""Render the scene and present it to the window.
|
||||
|
||||
Args:
|
||||
model: The MuJoCo model provided by the simulation.
|
||||
data: The MuJoCo data provided by the simulation.
|
||||
"""
|
||||
self._sync_renderer(model)
|
||||
self._viewer.Present(
|
||||
model,
|
||||
data,
|
||||
self.perturb,
|
||||
self.camera,
|
||||
self.vis_options,
|
||||
self.render_flags.flags,
|
||||
)
|
||||
|
||||
def stop(self) -> None:
|
||||
"""Stop the viewer."""
|
||||
self._is_running = False
|
||||
|
||||
def get_drop_file(self) -> str:
|
||||
"""Returns the path of the file dropped into the window, or empty string."""
|
||||
return self._viewer.GetDropFile()
|
||||
|
||||
def upload_image(
|
||||
self, tex_id: int, img: str | bytes, width: int, height: int, bpp: int
|
||||
) -> int:
|
||||
"""Uploads an image to the backend for GUI rendering.
|
||||
|
||||
The ID can be used in subsequent calls to update the texture data. An empty
|
||||
`img` argument will free the texture if it exists. A `tex_id` of 0 will
|
||||
create a new texture.
|
||||
|
||||
Args:
|
||||
tex_id: The texture ID.
|
||||
img: The image data as string or bytes.
|
||||
width: Width of the image.
|
||||
height: Height of the image.
|
||||
bpp: Bytes per pixel.
|
||||
|
||||
Returns:
|
||||
The texture ID.
|
||||
"""
|
||||
return self._viewer.UploadImage(tex_id, img, width, height, bpp)
|
||||
|
||||
def render_to_texture(
|
||||
self,
|
||||
model: mujoco.MjModel,
|
||||
data: mujoco.MjData,
|
||||
tex_id: int,
|
||||
width: int,
|
||||
height: int,
|
||||
) -> int:
|
||||
"""Renders the scene to a texture.
|
||||
|
||||
This function renders the scene from the current camera view into a texture.
|
||||
It handles buffer allocation internally.
|
||||
|
||||
Args:
|
||||
model: The MuJoCo model provided by the simulation.
|
||||
data: The MuJoCo data provided by the simulation.
|
||||
tex_id: The texture ID to render into (0 to create a new one).
|
||||
width: Width of the texture.
|
||||
height: Height of the texture.
|
||||
|
||||
Returns:
|
||||
The texture ID.
|
||||
"""
|
||||
self._sync_renderer(model)
|
||||
return self._viewer.RenderToTexture(
|
||||
model,
|
||||
data,
|
||||
self.camera,
|
||||
width,
|
||||
height,
|
||||
tex_id,
|
||||
)
|
||||
@@ -0,0 +1,46 @@
|
||||
// Copyright 2026 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
|
||||
//
|
||||
// https://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 <string_view>
|
||||
|
||||
#include <mujoco/mujoco.h>
|
||||
#include "third_party/mujoco/src/experimental/platform/sim/model_holder.h"
|
||||
#include "structs.h"
|
||||
#include <pybind11/pybind11.h>
|
||||
|
||||
namespace mujoco::python {
|
||||
|
||||
// Loads, parses, and compiles a MuJoCo model from the given file. Returns the
|
||||
// python mjData object (which also contains the compiled mjModel).
|
||||
py::object Parse(std::string_view filepath) {
|
||||
auto holder = platform::ModelHolder::FromFile(filepath);
|
||||
if (!holder->ok()) {
|
||||
throw py::value_error(
|
||||
std::string("Failed to load model from '") +
|
||||
std::string(filepath) + "': " + std::string(holder->error()));
|
||||
}
|
||||
mjModel* model = holder->ReleaseModel();
|
||||
mjData* data = holder->ReleaseData();
|
||||
py::object py_model = py::cast(MjModelWrapper(model));
|
||||
py::object py_data =
|
||||
py::cast(MjDataWrapper(py::cast<MjModelWrapper*>(py_model), data));
|
||||
return py_data;
|
||||
}
|
||||
|
||||
} // namespace mujoco::python
|
||||
|
||||
PYBIND11_MODULE(parser, m) {
|
||||
m.def("parse", &mujoco::python::Parse,
|
||||
pybind11::return_value_policy::take_ownership);
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
// Copyright 2026 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
|
||||
//
|
||||
// https://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 "third_party/mujoco/src/experimental/platform/hal/renderer.h"
|
||||
|
||||
#include <cstddef>
|
||||
#include <memory>
|
||||
#include <optional>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include <mujoco/mujoco.h>
|
||||
#include "third_party/mujoco/src/experimental/platform/hal/graphics_mode.h"
|
||||
#include "structs.h"
|
||||
#include <pybind11/eval.h>
|
||||
#include <pybind11/pybind11.h>
|
||||
#include <pybind11/pytypes.h>
|
||||
#include <pybind11/stl.h>
|
||||
|
||||
namespace mujoco::python {
|
||||
|
||||
class Renderer {
|
||||
public:
|
||||
using RendererImpl = mujoco::platform::Renderer;
|
||||
using GraphicsMode = mujoco::platform::GraphicsMode;
|
||||
|
||||
Renderer(const std::string& graphics_mode_str) {
|
||||
const GraphicsMode mode = mujoco::platform::GraphicsModeFromString(
|
||||
graphics_mode_str, GraphicsMode::FilamentOpenGl);
|
||||
impl_ = std::make_unique<RendererImpl>(nullptr, mode);
|
||||
}
|
||||
|
||||
void Init(const MjModelWrapper& model) { impl_->Init(model.get()); }
|
||||
|
||||
pybind11::bytes Render(const MjModelWrapper& model, MjDataWrapper& data,
|
||||
std::optional<MjvPerturbWrapper>& perturb,
|
||||
std::optional<MjvCameraWrapper>& camera,
|
||||
std::optional<MjvOptionWrapper>& vis_option, int width,
|
||||
int height) {
|
||||
std::vector<std::byte> pixels(width * height * 3);
|
||||
impl_->Render(
|
||||
model.get(), data.get(), perturb ? perturb.value().get() : nullptr,
|
||||
camera ? camera.value().get() : nullptr,
|
||||
vis_option ? vis_option.value().get() : nullptr, width, height, pixels);
|
||||
return pybind11::bytes((const char*)pixels.data(), pixels.size());
|
||||
}
|
||||
|
||||
pybind11::memoryview GetRenderFlags() {
|
||||
return pybind11::memoryview::from_buffer(
|
||||
impl_->GetRenderFlags(), {static_cast<pybind11::ssize_t>(mjNRNDFLAG)},
|
||||
{sizeof(mjtByte)});
|
||||
}
|
||||
|
||||
private:
|
||||
std::unique_ptr<RendererImpl> impl_;
|
||||
};
|
||||
|
||||
} // namespace mujoco::python
|
||||
|
||||
PYBIND11_MODULE(renderer, m) {
|
||||
pybind11::class_<mujoco::python::Renderer>(m, "Renderer")
|
||||
.def(pybind11::init<const std::string&>())
|
||||
.def("Init", &mujoco::python::Renderer::Init)
|
||||
.def("Render", &mujoco::python::Renderer::Render)
|
||||
.def("get_render_flags", &mujoco::python::Renderer::GetRenderFlags,
|
||||
pybind11::keep_alive<0, 1>());
|
||||
}
|
||||
@@ -0,0 +1,231 @@
|
||||
# Copyright 2026 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
|
||||
#
|
||||
# https://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.
|
||||
"""This script runs a simulation and viewer in separate processes communicating asynchronously.
|
||||
|
||||
In this example, we will run the viewer in an independent process communicating
|
||||
via multiprocessing queues. Controls are provided to simulate network transit
|
||||
latency and adjust the communication rates.
|
||||
|
||||
You must provide a mjcf model file via the first command-line argument.
|
||||
"""
|
||||
|
||||
import dataclasses
|
||||
import multiprocessing
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
|
||||
from absl import app as absl_app
|
||||
from absl import flags as absl_flags
|
||||
import mujoco
|
||||
from mujoco.experimental.studio import native_viewer as _viewer
|
||||
from mujoco.experimental.studio import sim as _sim
|
||||
from mujoco.experimental.studio import studio_app
|
||||
from mujoco.experimental.studio import ux
|
||||
import numpy as np
|
||||
|
||||
from mujoco.experimental.dear_imgui import dear_imgui as imgui
|
||||
|
||||
_GFX = absl_flags.DEFINE_string('gfx', '', 'Rendering graphics mode.')
|
||||
_WIDTH = absl_flags.DEFINE_integer('width', 1200, 'Width of the output image.')
|
||||
_HEIGHT = absl_flags.DEFINE_integer('height', 800, 'Height of the output image')
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class SimToView:
|
||||
"""A message sent from the simulation process to the viewer process."""
|
||||
|
||||
model: mujoco.MjModel | None = None
|
||||
data: mujoco.MjData | None = None
|
||||
state: np.ndarray | None = None
|
||||
state_sig: int = 0
|
||||
send_time: float = 0.0
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class ViewToSim:
|
||||
"""A message sent from the viewer process to the simulation process."""
|
||||
|
||||
state: np.ndarray | None = None
|
||||
state_sig: int = 0
|
||||
reset: bool = False
|
||||
send_rate: float = 60.0
|
||||
|
||||
|
||||
class Network:
|
||||
"""Simulated networking parameters."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.transit_buffer = []
|
||||
self.send_rate = 60.0
|
||||
self.network_delay = 0.2
|
||||
|
||||
def get_arrived(self, q: multiprocessing.Queue) -> SimToView | None:
|
||||
now = time.time()
|
||||
while not q.empty():
|
||||
self.transit_buffer.append(q.get())
|
||||
arrived = None
|
||||
while (
|
||||
self.transit_buffer
|
||||
and now >= self.transit_buffer[0].send_time + self.network_delay
|
||||
):
|
||||
arrived = self.transit_buffer.pop(0)
|
||||
return arrived
|
||||
|
||||
|
||||
def view(
|
||||
sim_to_view: multiprocessing.Queue,
|
||||
view_to_sim: multiprocessing.Queue,
|
||||
) -> None:
|
||||
"""Entry-point for process that renders the simulation."""
|
||||
# Block until the first message (containing the model) arrives.
|
||||
msg = sim_to_view.get()
|
||||
assert msg.model is not None, 'First message must contain the MuJoCo model.'
|
||||
|
||||
title = os.path.basename(sys.argv[0])
|
||||
xfrc_sig = int(mujoco.mjtState.mjSTATE_XFRC_APPLIED)
|
||||
xfrc_size = mujoco.mj_stateSize(msg.model, xfrc_sig)
|
||||
xfrc_state = np.zeros(xfrc_size, np.float64)
|
||||
|
||||
app = studio_app.StudioApp(msg.model, msg.data)
|
||||
network = Network()
|
||||
viewer = _viewer.NativeViewer(
|
||||
app.model,
|
||||
title=title,
|
||||
width=_WIDTH.value,
|
||||
height=_HEIGHT.value,
|
||||
gfx=_GFX.value,
|
||||
)
|
||||
|
||||
while viewer.is_running() and app.is_running():
|
||||
|
||||
# Determine which messages have arrived through the simulated network.
|
||||
arrived = network.get_arrived(sim_to_view)
|
||||
|
||||
# Update the camera and compute the perturbation.
|
||||
app.handle_mouse_events(viewer.camera, viewer.vis_options, viewer.perturb)
|
||||
|
||||
# Sync state from the backend if a new payload actually arrived.
|
||||
if arrived is not None and arrived.state is not None:
|
||||
mujoco.mj_setState(app.model, app.data, arrived.state, arrived.state_sig)
|
||||
mujoco.mj_forward(app.model, app.data)
|
||||
|
||||
# Always apply the perturbation forces from the viewer.
|
||||
app.apply_perturb(viewer.perturb)
|
||||
|
||||
# Transmit user interaction when we get a new state
|
||||
if arrived is not None:
|
||||
mujoco.mj_getState(app.model, app.data, xfrc_state, xfrc_sig)
|
||||
view_to_sim.put(
|
||||
ViewToSim(
|
||||
send_rate=network.send_rate, state=xfrc_state, state_sig=xfrc_sig
|
||||
)
|
||||
)
|
||||
|
||||
# Build the UI.
|
||||
ux.setup_theme(app.theme)
|
||||
if imgui.Begin(
|
||||
'Settings',
|
||||
flags=int(imgui.WindowFlags.AlwaysAutoResize)
|
||||
| int(imgui.WindowFlags.NoTitleBar)
|
||||
| int(imgui.WindowFlags.NoCollapse),
|
||||
):
|
||||
imgui.PushItemWidth(200.0)
|
||||
_, network.network_delay = imgui.SliderFloat(
|
||||
'Network Latency (s)', network.network_delay, 0.0, 2.0
|
||||
)
|
||||
updated, network.send_rate = imgui.SliderFloat(
|
||||
'Send Rate (Hz)', network.send_rate, 1.0, 120.0
|
||||
)
|
||||
if updated:
|
||||
view_to_sim.put(ViewToSim(send_rate=network.send_rate))
|
||||
|
||||
imgui.SetNextItemWidth(-1)
|
||||
if imgui.Button('Reset Simulation'):
|
||||
view_to_sim.put(ViewToSim(reset=True, send_rate=network.send_rate))
|
||||
|
||||
imgui.PopItemWidth()
|
||||
imgui.End()
|
||||
|
||||
viewer.sync(app.model, app.data)
|
||||
|
||||
|
||||
def sim(
|
||||
data: mujoco.MjData,
|
||||
model: mujoco.MjModel,
|
||||
sim_to_view: multiprocessing.Queue,
|
||||
view_to_sim: multiprocessing.Queue,
|
||||
view_process: multiprocessing.Process,
|
||||
) -> None:
|
||||
"""Entry-point for process that runs the simulation."""
|
||||
|
||||
sim_to_view.put(SimToView(model=model, data=data))
|
||||
|
||||
step_control = _sim.StepControl()
|
||||
integration_sig = int(mujoco.mjtState.mjSTATE_INTEGRATION)
|
||||
integration_size = mujoco.mj_stateSize(model, integration_sig)
|
||||
integration_state = np.empty(integration_size, np.float64)
|
||||
|
||||
msg = ViewToSim()
|
||||
last_send_time = time.time()
|
||||
|
||||
while view_process.is_alive():
|
||||
while not view_to_sim.empty():
|
||||
msg = view_to_sim.get()
|
||||
|
||||
if msg.reset:
|
||||
mujoco.mj_resetData(model, data)
|
||||
mujoco.mj_forward(model, data)
|
||||
msg.reset = False
|
||||
|
||||
# Apply perturbation forces received from the viewer process.
|
||||
if msg.state is not None:
|
||||
mujoco.mj_setState(model, data, msg.state, msg.state_sig)
|
||||
|
||||
# Advance the simulation keeping up with real-time.
|
||||
step_control.advance(model, data)
|
||||
|
||||
# Send the simulation state paced by the requested send_rate.
|
||||
now = time.time()
|
||||
if now - last_send_time >= 1.0 / max(1.0, msg.send_rate):
|
||||
mujoco.mj_getState(model, data, integration_state, integration_sig)
|
||||
sim_to_view.put(
|
||||
SimToView(
|
||||
state=integration_state,
|
||||
state_sig=integration_sig,
|
||||
send_time=now,
|
||||
)
|
||||
)
|
||||
last_send_time = now
|
||||
|
||||
|
||||
def main(argv: list[str]) -> None:
|
||||
app = studio_app.StudioApp.from_argv(argv)
|
||||
|
||||
# Queues for communication between the simulation and viewer processes.
|
||||
sim_to_view = multiprocessing.Queue()
|
||||
view_to_sim = multiprocessing.Queue()
|
||||
|
||||
# Start the viewer process.
|
||||
view_process = multiprocessing.Process(
|
||||
target=view, args=(sim_to_view, view_to_sim)
|
||||
)
|
||||
view_process.start()
|
||||
|
||||
# Start the simulation in the main process.
|
||||
sim(app.data, app.model, sim_to_view, view_to_sim, view_process)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
absl_app.run(main)
|
||||
@@ -0,0 +1,199 @@
|
||||
# Copyright 2026 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
|
||||
#
|
||||
# https://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.
|
||||
"""Example to run studio in the native viewer with responsive ImPlot UI.
|
||||
|
||||
This script runs a Studio viewer in-process and adds an 'Inspect Body' window
|
||||
using ImGui and ImPlot bindings to visualize selected body data. The example
|
||||
demonstrates how responsive UI layout rules are easily implemented.
|
||||
|
||||
Provide an MJCF model file via the first command-line argument to launch.
|
||||
"""
|
||||
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
|
||||
from absl import app as absl_app
|
||||
from absl import flags as absl_flags
|
||||
import mujoco
|
||||
from mujoco.experimental.studio import native_viewer as _viewer
|
||||
from mujoco.experimental.studio import studio_app
|
||||
import numpy as np
|
||||
|
||||
from mujoco.experimental.dear_imgui import dear_imgui as imgui
|
||||
from mujoco.experimental.implot import implot
|
||||
|
||||
_GFX = absl_flags.DEFINE_string('gfx', '', 'Rendering graphics mode.')
|
||||
_WIDTH = absl_flags.DEFINE_integer('width', 1200, 'Width of the output image.')
|
||||
_HEIGHT = absl_flags.DEFINE_integer('height', 800, 'Height of the output image')
|
||||
|
||||
|
||||
_N_HISTORY = 100
|
||||
|
||||
_PLOT_FLAGS = (
|
||||
implot.Flags.NoInputs.value # Disable pan/zoom mouse interaction.
|
||||
| implot.Flags.NoMenus.value # Disable right-click context menu.
|
||||
| implot.Flags.NoBoxSelect.value # Disable drag-to-select regions.
|
||||
)
|
||||
|
||||
_AXIS_FLAGS = (
|
||||
implot.AxisFlags.NoGridLines.value # Hide background grid lines.
|
||||
| implot.AxisFlags.NoTickMarks.value # Hide small tick marks on the axis.
|
||||
)
|
||||
|
||||
|
||||
def _setup_plot_flags(plot_size: imgui.Vec2) -> int:
|
||||
flags = _PLOT_FLAGS
|
||||
if min(plot_size.x, plot_size.y) < 300:
|
||||
flags |= implot.Flags.NoTitle.value
|
||||
if min(plot_size.x, plot_size.y) < 200:
|
||||
flags |= implot.Flags.NoLegend.value
|
||||
return flags
|
||||
|
||||
|
||||
def _setup_time_axis(plot_size: imgui.Vec2) -> None:
|
||||
flags = _AXIS_FLAGS
|
||||
if plot_size.x < 300:
|
||||
flags |= implot.AxisFlags.NoTickLabels.value
|
||||
implot.SetupAxis(implot.Axis.X1, '', flags)
|
||||
implot.SetupAxisLimits(implot.Axis.X1, 0, _N_HISTORY)
|
||||
|
||||
|
||||
def _setup_xpos_axis(centroid: list[np.ndarray], plot_size: imgui.Vec2) -> None:
|
||||
flags = _AXIS_FLAGS
|
||||
if plot_size.y < 300:
|
||||
flags |= implot.AxisFlags.NoTickLabels.value
|
||||
implot.SetupAxis(implot.Axis.Y1, '', flags)
|
||||
min_y = min(c[1] for c in centroid)
|
||||
max_y = max(c[1] for c in centroid)
|
||||
margin = max((max_y - min_y) * 0.1, 0.05)
|
||||
implot.SetupAxisLimits(
|
||||
implot.Axis.Y1,
|
||||
min_y - margin,
|
||||
max_y + margin,
|
||||
cond=implot.Cond.Always,
|
||||
)
|
||||
|
||||
|
||||
def _setup_angle_axis(plot_size: imgui.Vec2) -> None:
|
||||
flags = _AXIS_FLAGS
|
||||
if plot_size.y < 300:
|
||||
flags |= implot.AxisFlags.NoTickLabels.value
|
||||
implot.SetupAxis(implot.Axis.Y1, '', flags)
|
||||
implot.SetupAxisLimits(implot.Axis.Y1, -185.0, 185.0)
|
||||
implot.SetupAxisTicks(
|
||||
implot.Axis.Y1,
|
||||
[-180.0, -90.0, 0.0, 90.0, 180.0],
|
||||
['-180', '-90', '0', '90', '180'],
|
||||
)
|
||||
|
||||
|
||||
def main(argv: list[str]) -> None:
|
||||
app = studio_app.StudioApp.from_argv(argv)
|
||||
title = os.path.basename(sys.argv[0])
|
||||
|
||||
# Initialize the viewer.
|
||||
viewer = _viewer.NativeViewer(
|
||||
app.model,
|
||||
title=title,
|
||||
width=_WIDTH.value,
|
||||
height=_HEIGHT.value,
|
||||
gfx=_GFX.value,
|
||||
)
|
||||
|
||||
# Variables for the custom UI.
|
||||
centroid = [np.zeros(3) for _ in range(_N_HISTORY)]
|
||||
euler = [np.zeros(3) for _ in range(_N_HISTORY)]
|
||||
body_id = -1
|
||||
|
||||
# Main viewer loop.
|
||||
while viewer.is_running():
|
||||
if not app.update(viewer.camera, viewer.vis_options, viewer.perturb):
|
||||
break
|
||||
|
||||
# Build standard Studio UI.
|
||||
app.build_gui(viewer.camera, viewer.vis_options, viewer.render_flags)
|
||||
|
||||
# Inspect the perturb.select body
|
||||
if viewer.perturb.select > 0:
|
||||
body_id = viewer.perturb.select
|
||||
|
||||
# Display selected body information.
|
||||
if body_id > 0:
|
||||
body_name = mujoco.mj_id2name(
|
||||
app.model, int(mujoco.mjtObj.mjOBJ_BODY), body_id
|
||||
)
|
||||
|
||||
imgui.SetNextWindowSize(imgui.Vec2(1200, 600), imgui.Cond.FirstUseEver)
|
||||
|
||||
# Note: The window title uses the special "###" markup to ensure the imgui
|
||||
# ID for the window is constant for all body names. This is needed for
|
||||
# the window to retain its state for all bodies.
|
||||
window_title = f'Inspect Body {body_name or "(???)"!r} ({body_id})###Plot'
|
||||
if imgui.Begin(window_title):
|
||||
avail = imgui.GetContentRegionAvail()
|
||||
wide = avail.x > avail.y
|
||||
|
||||
# Add a small padding factor to prevent scrollbars.
|
||||
plot_size = imgui.Vec2(
|
||||
avail.x * 0.5 - 4 if wide else avail.x,
|
||||
avail.y if wide else avail.y * 0.5 - 4,
|
||||
)
|
||||
|
||||
plot_flags = _setup_plot_flags(plot_size)
|
||||
if implot.BeginPlot('Centroid vs Time', plot_size, flags=plot_flags):
|
||||
_setup_time_axis(plot_size)
|
||||
_setup_xpos_axis(centroid, plot_size)
|
||||
implot.PlotLine('x', range(_N_HISTORY), [c[0] for c in centroid])
|
||||
implot.PlotLine('y', range(_N_HISTORY), [c[1] for c in centroid])
|
||||
implot.PlotLine('z', range(_N_HISTORY), [c[2] for c in centroid])
|
||||
implot.EndPlot()
|
||||
|
||||
if wide:
|
||||
imgui.SameLine()
|
||||
|
||||
if implot.BeginPlot('Euler Angle vs Time', plot_size, flags=plot_flags):
|
||||
_setup_time_axis(plot_size)
|
||||
_setup_angle_axis(plot_size)
|
||||
implot.PlotLine('roll', range(_N_HISTORY), [e[0] for e in euler])
|
||||
implot.PlotLine('pitch', range(_N_HISTORY), [e[1] for e in euler])
|
||||
implot.PlotLine('yaw', range(_N_HISTORY), [e[2] for e in euler])
|
||||
implot.EndPlot()
|
||||
imgui.End()
|
||||
|
||||
# Update plot data
|
||||
centroid.pop(0)
|
||||
euler.pop(0)
|
||||
if body_id > 0:
|
||||
centroid.append(app.data.xpos[body_id].copy())
|
||||
# Convert quaternion to Euler angles via rotation matrix.
|
||||
quat = app.data.xquat[body_id]
|
||||
mat = np.zeros(9)
|
||||
mujoco.mju_quat2Mat(mat, quat)
|
||||
# mat is row-major 3x3: R[i,j] = mat[3*i + j].
|
||||
roll = math.atan2(mat[7], mat[8])
|
||||
pitch = math.atan2(-mat[6], math.sqrt(mat[7] ** 2 + mat[8] ** 2))
|
||||
yaw = math.atan2(mat[3], mat[0])
|
||||
euler.append(np.degrees(np.array([roll, pitch, yaw])))
|
||||
else:
|
||||
centroid.append(np.zeros(3))
|
||||
euler.append(np.zeros(3))
|
||||
|
||||
viewer.sync(app.model, app.data)
|
||||
|
||||
viewer.stop()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
absl_app.run(main)
|
||||
@@ -0,0 +1,73 @@
|
||||
# Copyright 2026 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
|
||||
#
|
||||
# https://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.
|
||||
"""Render a MuJoCo model to an image."""
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
from absl import app
|
||||
from absl import flags
|
||||
import mujoco
|
||||
from mujoco.experimental.studio import parser
|
||||
from mujoco.experimental.studio import renderer
|
||||
from PIL import Image
|
||||
|
||||
_MODEL = flags.DEFINE_string('model', '', 'Model file to load.')
|
||||
_OUTPUT = flags.DEFINE_string('output', '', 'Output file to save.')
|
||||
_GFX = flags.DEFINE_string('gfx', '', 'Renderer to use.')
|
||||
_WIDTH = flags.DEFINE_integer('width', 320, 'Width of the output image.')
|
||||
_HEIGHT = flags.DEFINE_integer('height', 240, 'Height of the output image.')
|
||||
_STEPS = flags.DEFINE_integer('steps', 1, 'Number of steps before render.')
|
||||
|
||||
|
||||
def main(argv):
|
||||
if len(argv) > 1:
|
||||
raise app.UsageError('Too many command-line arguments.')
|
||||
if not _MODEL.value:
|
||||
raise ValueError('`model` flag is required.')
|
||||
if not _OUTPUT.value:
|
||||
raise ValueError('`output flag is required.')
|
||||
|
||||
try:
|
||||
data = parser.parse(_MODEL.value)
|
||||
model = data.model
|
||||
except Exception as ex: # pylint: disable=broad-except
|
||||
print(f'Error loading model from `{_MODEL.value}`: {ex}')
|
||||
sys.exit(-1)
|
||||
|
||||
for _ in range(_STEPS.value):
|
||||
mujoco.mj_step(model, data)
|
||||
|
||||
try:
|
||||
r = renderer.Renderer(_GFX.value)
|
||||
r.Init(model)
|
||||
pixels = r.Render(
|
||||
model, data, None, None, None, _WIDTH.value, _HEIGHT.value
|
||||
)
|
||||
except Exception as ex: # pylint: disable=broad-except
|
||||
print(f'Error rendering model: {ex}')
|
||||
sys.exit(-2)
|
||||
|
||||
try:
|
||||
img = Image.frombytes('RGB', (_WIDTH.value, _HEIGHT.value), pixels)
|
||||
img.save(_OUTPUT.value, format=os.path.splitext(_OUTPUT.value)[1][1:])
|
||||
except Exception as ex: # pylint: disable=broad-except
|
||||
print(f'Error saving image to `{_OUTPUT.value}`: {ex}')
|
||||
sys.exit(-3)
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
app.run(main)
|
||||
@@ -0,0 +1,65 @@
|
||||
// Copyright 2026 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
|
||||
//
|
||||
// https://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.
|
||||
|
||||
// Python bindings for MuJoCo platform simulation components.
|
||||
|
||||
#include <mujoco/mujoco.h>
|
||||
#include "third_party/mujoco/src/experimental/platform/sim/step_control.h"
|
||||
#include "structs.h"
|
||||
#include <pybind11/pybind11.h>
|
||||
|
||||
namespace py = pybind11;
|
||||
|
||||
using StepControl = mujoco::platform::StepControl;
|
||||
|
||||
PYBIND11_MODULE(sim, m) {
|
||||
m.doc() = "MuJoCo platform simulation bindings for Link.";
|
||||
|
||||
py::enum_<StepControl::Status>(m, "StepStatus")
|
||||
.value("OK", StepControl::Status::kOk)
|
||||
.value("PAUSED", StepControl::Status::kPaused)
|
||||
.value("VISCOUS_PAUSED", StepControl::Status::kViscousPaused)
|
||||
.value("AUTO_RESET", StepControl::Status::kAutoReset)
|
||||
.value("DIVERGED", StepControl::Status::kDiverged);
|
||||
|
||||
py::enum_<StepControl::PauseState>(m, "PauseState")
|
||||
.value("UNPAUSED", StepControl::PauseState::kUnpaused)
|
||||
.value("NORMAL_PAUSED", StepControl::PauseState::kNormalPaused)
|
||||
.value("VISCOUS_PAUSED", StepControl::PauseState::kViscousPaused);
|
||||
|
||||
py::class_<StepControl>(m, "StepControl")
|
||||
.def(py::init<>())
|
||||
.def(
|
||||
"advance",
|
||||
[](StepControl& self, mujoco::python::MjModelWrapper& model,
|
||||
mujoco::python::MjDataWrapper& data) {
|
||||
return self.Advance(model.get(), data.get());
|
||||
},
|
||||
py::arg("model"), py::arg("data"),
|
||||
"Step physics forward, respecting speed settings and refresh budget.")
|
||||
.def("force_sync", &StepControl::ForceSync,
|
||||
"Ensures the next Advance() will synchronize time and step once.")
|
||||
.def("get_speed", &StepControl::GetSpeed,
|
||||
"Returns the desired simulation speed as a percentage of real time.")
|
||||
.def("get_speed_measured", &StepControl::GetSpeedMeasured,
|
||||
"Returns the measured simulation speed.")
|
||||
.def("set_speed", &StepControl::SetSpeed, py::arg("speed"),
|
||||
"Sets the desired speed (clamped to [0.1%, 100%]).")
|
||||
.def("set_pause_state", &StepControl::SetPauseState, py::arg("state"),
|
||||
"Sets the pause state of the simulation.")
|
||||
.def("get_pause_state", &StepControl::GetPauseState,
|
||||
"Returns the current pause state.")
|
||||
.def("request_single_step", &StepControl::RequestSingleStep,
|
||||
"Request a single step if paused.");
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
# Copyright 2026 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
|
||||
#
|
||||
# https://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.
|
||||
"""This script runs Studio from Python, visualized in a native viewer."""
|
||||
|
||||
from absl import app as absl_app
|
||||
from absl import flags as absl_flags
|
||||
from mujoco.experimental.studio import native_viewer
|
||||
from mujoco.experimental.studio import studio_app
|
||||
|
||||
_GFX = absl_flags.DEFINE_string('gfx', '', 'Rendering graphics mode.')
|
||||
_WIDTH = absl_flags.DEFINE_integer('width', 1200, 'Width of the output image.')
|
||||
_HEIGHT = absl_flags.DEFINE_integer('height', 800, 'Height of the output image')
|
||||
|
||||
|
||||
def main(argv: list[str]) -> None:
|
||||
app = studio_app.StudioApp.from_argv(argv)
|
||||
|
||||
# Initialize the viewer.
|
||||
viewer = native_viewer.NativeViewer(
|
||||
app.model,
|
||||
width=_WIDTH.value,
|
||||
height=_HEIGHT.value,
|
||||
gfx=_GFX.value,
|
||||
)
|
||||
|
||||
# Main viewer loop.
|
||||
while viewer.is_running():
|
||||
if not app.update_from_viewer(viewer):
|
||||
break
|
||||
|
||||
app.build_gui(viewer.camera, viewer.vis_options, viewer.render_flags)
|
||||
|
||||
viewer.sync(app.model, app.data)
|
||||
|
||||
viewer.stop()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
absl_app.run(main)
|
||||
@@ -0,0 +1,440 @@
|
||||
# Copyright 2026 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
|
||||
#
|
||||
# https://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.
|
||||
"""Viewer-agnostic Python implementation of Studio.
|
||||
|
||||
Architecture:
|
||||
StudioApp owns the simulation state (model, data) and the UI logic.
|
||||
Viewer classes (e.g., NativeViewer) own the window (if required), renderer,
|
||||
camera, and visualization options. The viewer never stores references to
|
||||
model or data. Instead, the caller passes them each frame via
|
||||
viewer.sync(model, data). This ensures the viewer always renders the current
|
||||
model, even if StudioApp.load_model_from_file() swaps it.
|
||||
|
||||
The class can be used to implement the full Studio application in Python. By
|
||||
using the more granular member functions it can also build simple apps that only
|
||||
use a subset of the Studio UI. This configuration is fully dynamic, there is
|
||||
nothing to configure in advance, you can change your app by changing the
|
||||
functions that get called each frame. This class is also viewer-agnostic and as
|
||||
such does not own camera, vis_options or perturb objects (these are provided by
|
||||
the viewer).
|
||||
|
||||
See the sample/ folder for usage examples.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
import mujoco
|
||||
from mujoco.experimental.studio import parser
|
||||
from mujoco.experimental.studio import sim
|
||||
from mujoco.experimental.studio import studio_app_events as events
|
||||
from mujoco.experimental.studio import ux
|
||||
from mujoco.experimental.studio import viewer_protocol
|
||||
import numpy as np
|
||||
|
||||
from mujoco.experimental.dear_imgui import dear_imgui as imgui
|
||||
|
||||
|
||||
def load_model_from_file(
|
||||
model_path: str,
|
||||
) -> tuple[mujoco.MjModel, mujoco.MjData] | None:
|
||||
"""Loads a model and data from a file path."""
|
||||
try:
|
||||
data = parser.parse(model_path)
|
||||
return data.model, data
|
||||
except Exception as ex: # pylint: disable=broad-except
|
||||
print(f'Error loading model from {model_path!r}: {ex}')
|
||||
return None
|
||||
|
||||
|
||||
class StudioApp:
|
||||
"""Viewer-agnostic Python implementation of Studio."""
|
||||
|
||||
@classmethod
|
||||
def from_argv(cls, argv: list[str]) -> 'StudioApp':
|
||||
"""Constructs a StudioApp by parsing a model path from command-line args."""
|
||||
if len(argv) < 2:
|
||||
model = mujoco.MjSpec().compile()
|
||||
data = mujoco.MjData(model)
|
||||
app = cls(model, data)
|
||||
app.step_control.set_pause_state(sim.PauseState.NORMAL_PAUSED)
|
||||
return app
|
||||
|
||||
model_path = argv[1]
|
||||
|
||||
res = load_model_from_file(model_path)
|
||||
if res is None:
|
||||
sys.exit(-1)
|
||||
model, data = res
|
||||
|
||||
app = cls(model, data)
|
||||
app.model_path = model_path
|
||||
return app
|
||||
|
||||
def load_model_from_file(
|
||||
self, model_path: str
|
||||
) -> tuple[mujoco.MjModel, mujoco.MjData] | None:
|
||||
"""Loads a new model from a file, replacing the current model and data."""
|
||||
res = load_model_from_file(model_path)
|
||||
if res is None:
|
||||
self.status = f'Error loading model from {model_path!r}'
|
||||
return None
|
||||
|
||||
model, data = res
|
||||
self.model = model
|
||||
self.data = data
|
||||
self.model_path = model_path
|
||||
self.step_control = sim.StepControl()
|
||||
self.ux_state = ux.UxState()
|
||||
self.status = f'Loaded: {os.path.basename(model_path)!r}'
|
||||
|
||||
return model, data
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: mujoco.MjModel,
|
||||
data: mujoco.MjData,
|
||||
):
|
||||
"""Initializes the Studio application."""
|
||||
self.model = model
|
||||
self.data = data
|
||||
self.model_path = ''
|
||||
|
||||
self.step_control = sim.StepControl()
|
||||
self.ux_state = ux.UxState()
|
||||
self.theme = ux.GuiTheme.LIGHT
|
||||
self.show_stats = False
|
||||
self.show_solver = False
|
||||
self.should_quit = False
|
||||
self.status = 'Ready'
|
||||
# TODO(matijak): This should be part of the viewer, also making a struct to
|
||||
# pass it around with the camera would be convenient.
|
||||
self._cam_speed = 0.001
|
||||
|
||||
def handle_vis_options_keyboard_events(
|
||||
self,
|
||||
vis_options: mujoco.MjvOption,
|
||||
is_freecam_wasd: bool,
|
||||
) -> bool:
|
||||
"""Toggles visualization flags based on keyboard shortcuts.
|
||||
|
||||
Args:
|
||||
vis_options: The visualization options to modify.
|
||||
is_freecam_wasd: If True, keys Q/E/A/D are reserved for camera movement
|
||||
and will not toggle visualization flags.
|
||||
|
||||
Returns:
|
||||
True if a key was handled, False otherwise.
|
||||
"""
|
||||
if imgui.GetIO().WantCaptureKeyboard:
|
||||
return False
|
||||
|
||||
return events.handle_vis_options_keyboard_events(
|
||||
vis_options, is_freecam_wasd
|
||||
)
|
||||
|
||||
def handle_step_control_keyboard_events(self) -> bool:
|
||||
"""Handles keyboard shortcuts for simulation stepping control.
|
||||
|
||||
Returns:
|
||||
True if a key was handled, False otherwise.
|
||||
"""
|
||||
if imgui.GetIO().WantCaptureKeyboard:
|
||||
return False
|
||||
|
||||
return events.handle_step_control_keyboard_events(
|
||||
self.model, self.data, self.step_control, self.ux_state
|
||||
)
|
||||
|
||||
def handle_freecam_wasd_keyboard_events(
|
||||
self,
|
||||
camera: mujoco.MjvCamera,
|
||||
) -> bool:
|
||||
"""Handles keyboard shortcuts for free camera movement."""
|
||||
if imgui.GetIO().WantCaptureKeyboard:
|
||||
return False
|
||||
|
||||
handled, self._cam_speed = events.handle_freecam_wasd_keyboard_events(
|
||||
self.model, self.data, camera, self._cam_speed
|
||||
)
|
||||
return handled
|
||||
|
||||
def handle_keyboard_events(
|
||||
self,
|
||||
camera: mujoco.MjvCamera,
|
||||
vis_options: mujoco.MjvOption,
|
||||
) -> bool:
|
||||
"""Handle keyboard events according to Studio's bindings."""
|
||||
if imgui.GetIO().WantCaptureKeyboard:
|
||||
return False
|
||||
|
||||
is_freecam_wasd = self.ux_state.camera_index == ux.FREE_CAMERA_IDX
|
||||
if events.handle_step_control_keyboard_events(
|
||||
self.model, self.data, self.step_control, self.ux_state
|
||||
):
|
||||
return True
|
||||
|
||||
if events.handle_camera_select_keyboard_events(
|
||||
self.model, camera, self.ux_state
|
||||
):
|
||||
return True
|
||||
|
||||
if events.handle_vis_options_keyboard_events(vis_options, is_freecam_wasd):
|
||||
return True
|
||||
|
||||
if is_freecam_wasd:
|
||||
handled, self._cam_speed = events.handle_freecam_wasd_keyboard_events(
|
||||
self.model, self.data, camera, self._cam_speed
|
||||
)
|
||||
if handled:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def handle_camera_tracking_mouse_events(
|
||||
self,
|
||||
camera: mujoco.MjvCamera,
|
||||
vis_options: mujoco.MjvOption,
|
||||
) -> None:
|
||||
"""Handles mouse events for camera tracking."""
|
||||
if imgui.GetIO().WantCaptureMouse:
|
||||
return
|
||||
|
||||
events.handle_camera_tracking_mouse_events(
|
||||
self.model, self.data, camera, vis_options, self.ux_state
|
||||
)
|
||||
|
||||
def handle_mouse_events(
|
||||
self,
|
||||
camera: mujoco.MjvCamera,
|
||||
vis_options: mujoco.MjvOption,
|
||||
perturb: mujoco.MjvPerturb,
|
||||
) -> None:
|
||||
"""Handles mouse events."""
|
||||
if imgui.GetIO().WantCaptureMouse:
|
||||
return
|
||||
|
||||
events.handle_mouse_events(
|
||||
self.model, self.data, camera, vis_options, perturb, self.ux_state
|
||||
)
|
||||
|
||||
def reset_physics(self) -> None:
|
||||
"""Reset the physics."""
|
||||
mujoco.mj_resetData(self.model, self.data)
|
||||
mujoco.mj_forward(self.model, self.data)
|
||||
|
||||
def apply_perturb(self, perturb: mujoco.MjvPerturb) -> None:
|
||||
"""Apply perturbation the model."""
|
||||
if self.step_control.get_pause_state() != sim.PauseState.NORMAL_PAUSED:
|
||||
sig = mujoco.mjtState.mjSTATE_XFRC_APPLIED.value
|
||||
size = mujoco.mj_stateSize(self.model, sig)
|
||||
zero_state = np.zeros(size, np.float64)
|
||||
mujoco.mj_setState(self.model, self.data, zero_state, sig)
|
||||
mujoco.mjv_applyPerturbPose(self.model, self.data, perturb, 0)
|
||||
mujoco.mjv_applyPerturbForce(self.model, self.data, perturb)
|
||||
else:
|
||||
mujoco.mjv_applyPerturbPose(self.model, self.data, perturb, 1)
|
||||
|
||||
def update_physics(self, perturb: mujoco.MjvPerturb) -> None:
|
||||
"""Applies the purturbations and advances the physics."""
|
||||
self.apply_perturb(perturb)
|
||||
|
||||
advance_status = self.step_control.advance(self.model, self.data)
|
||||
if advance_status == sim.StepStatus.AUTO_RESET:
|
||||
self.reset_physics()
|
||||
|
||||
def reset_physics_gui(self) -> None:
|
||||
"""GUI to Reset the physics i.e., the reset button."""
|
||||
button_size = imgui.GetFrameHeight()
|
||||
square_size = imgui.Vec2(button_size, button_size)
|
||||
icon_reset_model = '\uf0e2' # FontAwesome "undo" icon.
|
||||
if imgui.Button(icon_reset_model, square_size):
|
||||
self.reset_physics()
|
||||
imgui.SetItemTooltip('Reset')
|
||||
|
||||
def is_running(self) -> bool:
|
||||
"""Returns True if the application should continue running (called by update())."""
|
||||
return not self.should_quit
|
||||
|
||||
def update(
|
||||
self,
|
||||
camera: mujoco.MjvCamera,
|
||||
vis_options: mujoco.MjvOption,
|
||||
perturb: mujoco.MjvPerturb,
|
||||
drop_file: str = '',
|
||||
) -> bool:
|
||||
"""Update the simulation and handle user input.
|
||||
|
||||
Handles mouse input to compute perturbations or camera motion.
|
||||
Handles keyboard input e.g., for keybindings or camera motion.
|
||||
Applies the purturbations and advances the physics.
|
||||
The argument objects are provided by the viewer.
|
||||
|
||||
Args:
|
||||
camera: The MuJoCo camera object.
|
||||
vis_options: The MuJoCo visualization options.
|
||||
perturb: The MuJoCo perturbation object.
|
||||
drop_file: Path of a file dropped into the viewer window. If non-empty the
|
||||
current model is replaced with the dropped file.
|
||||
|
||||
Returns:
|
||||
Whether the application should continue running, this is a
|
||||
convenience to allow this function to be used in a while loop.
|
||||
"""
|
||||
if drop_file:
|
||||
self.load_model_from_file(drop_file)
|
||||
|
||||
self.handle_mouse_events(camera, vis_options, perturb)
|
||||
self.handle_keyboard_events(camera, vis_options)
|
||||
self.update_physics(perturb)
|
||||
return self.is_running()
|
||||
|
||||
def update_from_viewer(self, viewer: viewer_protocol.Viewer) -> bool:
|
||||
"""Convenience wrapper around update() that unpacks viewer attributes."""
|
||||
return self.update(
|
||||
viewer.camera,
|
||||
viewer.vis_options,
|
||||
viewer.perturb,
|
||||
drop_file=viewer.get_drop_file(),
|
||||
)
|
||||
|
||||
def build_gui(
|
||||
self,
|
||||
camera: mujoco.MjvCamera,
|
||||
vis_options: mujoco.MjvOption,
|
||||
render_flags: ux.RenderFlags,
|
||||
) -> None:
|
||||
"""Emit full Studio UI."""
|
||||
ux.setup_theme(self.theme)
|
||||
ux.configure_docking_layout()
|
||||
|
||||
# -- Main menu bar --------------------------------------------------------
|
||||
if imgui.BeginMainMenuBar():
|
||||
if imgui.BeginMenu('File'):
|
||||
if imgui.MenuItem('Quit'):
|
||||
self.should_quit = True
|
||||
imgui.EndMenu()
|
||||
if imgui.BeginMenu('Simulation'):
|
||||
imgui.EndMenu()
|
||||
if imgui.BeginMenu('Charts'):
|
||||
if imgui.MenuItem('Solver', '', self.show_solver):
|
||||
self.show_solver = not self.show_solver
|
||||
if imgui.MenuItem('Stats', '', self.show_stats):
|
||||
self.show_stats = not self.show_stats
|
||||
imgui.EndMenu()
|
||||
if imgui.BeginMenu('Help'):
|
||||
if imgui.MenuItem('Stats', '', self.show_stats):
|
||||
self.show_stats = not self.show_stats
|
||||
imgui.Separator()
|
||||
version = f'Version {mujoco.mj_versionString()}'
|
||||
imgui.MenuItem(version)
|
||||
imgui.EndMenu()
|
||||
imgui.EndMainMenuBar()
|
||||
|
||||
# -- Tool Bar -------------------------------------------------------------
|
||||
if imgui.Begin('ToolBar'):
|
||||
imgui.PushStyleVar(imgui.StyleVar.CellPadding, imgui.Vec2(0, 0))
|
||||
if imgui.BeginTable('##ToolBarTable', 2):
|
||||
imgui.TableSetupColumn('', int(imgui.TableColumnFlags.WidthStretch))
|
||||
imgui.TableSetupColumn('', int(imgui.TableColumnFlags.WidthFixed))
|
||||
|
||||
imgui.TableNextColumn()
|
||||
self.reset_physics_gui()
|
||||
|
||||
imgui.SameLine()
|
||||
ux.step_control_gui(self.model, self.step_control, self.ux_state)
|
||||
|
||||
imgui.TableNextColumn()
|
||||
ux.camera_selection_gui(self.model, self.data, camera, self.ux_state)
|
||||
|
||||
imgui.SameLine()
|
||||
ux.label_selection_gui(vis_options)
|
||||
|
||||
imgui.SameLine()
|
||||
ux.frame_selection_gui(vis_options)
|
||||
|
||||
imgui.SameLine()
|
||||
changed, self.theme = ux.theme_select_gui(self.theme)
|
||||
if changed:
|
||||
ux.setup_theme(self.theme)
|
||||
|
||||
imgui.EndTable()
|
||||
imgui.PopStyleVar()
|
||||
imgui.End()
|
||||
|
||||
# -- Left pane: Options ---------------------------------------------------
|
||||
node_flags = int(imgui.TreeNodeFlags.SpanAvailWidth) | int(
|
||||
imgui.TreeNodeFlags.Framed
|
||||
)
|
||||
|
||||
imgui.Begin('Options')
|
||||
if imgui.TreeNodeEx('Physics Settings', node_flags):
|
||||
ux.physics_gui(self.model)
|
||||
imgui.TreePop()
|
||||
if imgui.TreeNodeEx('Rendering Settings', node_flags):
|
||||
ux.rendering_gui(self.model, vis_options, render_flags)
|
||||
imgui.TreePop()
|
||||
if imgui.TreeNodeEx('Visibility Groups', node_flags):
|
||||
ux.groups_gui(self.model, vis_options)
|
||||
imgui.TreePop()
|
||||
if imgui.TreeNodeEx('Visualization', node_flags):
|
||||
ux.visualization_gui(self.model, vis_options, camera)
|
||||
imgui.TreePop()
|
||||
imgui.End()
|
||||
|
||||
# -- Right pane: Inspector ------------------------------------------------
|
||||
imgui.Begin('Inspector')
|
||||
if imgui.TreeNodeEx('Noise', node_flags):
|
||||
ux.noise_gui(self.model, self.data, self.ux_state)
|
||||
imgui.TreePop()
|
||||
if imgui.TreeNodeEx('Joints', node_flags):
|
||||
ux.joints_gui(self.model, self.data, vis_options)
|
||||
imgui.TreePop()
|
||||
if imgui.TreeNodeEx('Controls', node_flags):
|
||||
ux.controls_gui(self.model, self.data, vis_options)
|
||||
imgui.TreePop()
|
||||
if imgui.TreeNodeEx(
|
||||
'Sensors', node_flags | int(imgui.TreeNodeFlags.DefaultOpen)
|
||||
):
|
||||
ux.sensor_gui(self.model, self.data)
|
||||
imgui.TreePop()
|
||||
if imgui.TreeNodeEx('Watch', node_flags):
|
||||
ux.watch_gui(self.model, self.data, self.ux_state)
|
||||
imgui.TreePop()
|
||||
if imgui.TreeNodeEx('State', node_flags):
|
||||
ux.state_gui(self.model, self.data, self.ux_state)
|
||||
imgui.TreePop()
|
||||
imgui.End()
|
||||
|
||||
# -- Floating windows -----------------------------------------------------
|
||||
if self.show_solver:
|
||||
_, self.show_solver = imgui.Begin('Solver', self.show_solver)
|
||||
ux.counts_gui(self.model, self.data)
|
||||
ux.convergence_gui(self.model, self.data)
|
||||
imgui.End()
|
||||
|
||||
if self.show_stats:
|
||||
_, self.show_stats = imgui.Begin('Stats', self.show_stats)
|
||||
paused = self.step_control.get_pause_state() != sim.PauseState.UNPAUSED
|
||||
ux.stats_gui(self.model, self.data, paused, 0.0)
|
||||
imgui.End()
|
||||
|
||||
# -- Status bar -----------------------------------------------------------
|
||||
imgui.PushStyleVar(imgui.StyleVar.CellPadding, imgui.Vec2(0, 0))
|
||||
imgui.PushStyleVar(imgui.StyleVar.FramePadding, imgui.Vec2(0, 0))
|
||||
imgui.PushStyleVar(imgui.StyleVar.WindowPadding, imgui.Vec2(0, 0))
|
||||
if imgui.Begin('StatusBar'):
|
||||
imgui.Text(self.status)
|
||||
imgui.End()
|
||||
imgui.PopStyleVar(3)
|
||||
@@ -0,0 +1,592 @@
|
||||
# Copyright 2026 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
|
||||
#
|
||||
# https://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.
|
||||
"""Temporary event handling functions for StudioApp."""
|
||||
|
||||
# TODO(matijak): These free functions implement the keyboard and mouse event
|
||||
# handling for Studio. They are separated from the main StudioApp class to keep
|
||||
# clarify the long-term API and avoid cluttering it with a large amount of
|
||||
# temporary code. When studio/platform has a proper API for registering key
|
||||
# bindings and mouse behaviour, the event handling functions will delegate to
|
||||
# code shared with the C++ studio application.
|
||||
|
||||
import mujoco
|
||||
from mujoco.experimental.studio import sim
|
||||
from mujoco.experimental.studio import ux
|
||||
import numpy as np
|
||||
|
||||
from mujoco.experimental.dear_imgui import dear_imgui as imgui
|
||||
|
||||
|
||||
def handle_vis_options_keyboard_events(
|
||||
vis_options: mujoco.MjvOption,
|
||||
is_freecam_wasd: bool,
|
||||
) -> bool:
|
||||
"""Toggles visualization flags based on keyboard shortcuts.
|
||||
|
||||
Args:
|
||||
vis_options: The visualization options to modify.
|
||||
is_freecam_wasd: If True, keys Q/E/A/D are reserved for camera movement and
|
||||
will not toggle visualization flags.
|
||||
|
||||
Returns:
|
||||
True if a key was handled, False otherwise.
|
||||
"""
|
||||
if imgui.GetIO().WantCaptureKeyboard:
|
||||
return False
|
||||
|
||||
pressed = imgui.IsKeyChordPressed
|
||||
|
||||
# Frame and label cycling.
|
||||
if pressed(imgui.Key.F6):
|
||||
vis_options.frame = (vis_options.frame + 1) % mujoco.mjtFrame.mjNFRAME.value
|
||||
elif pressed(imgui.Key.F7):
|
||||
vis_options.label = (vis_options.label + 1) % mujoco.mjtLabel.mjNLABEL.value
|
||||
|
||||
# Visualization flag toggles (single-key shortcuts).
|
||||
elif pressed(imgui.Key.H):
|
||||
vis_options.flags[mujoco.mjtVisFlag.mjVIS_CONVEXHULL] ^= 1
|
||||
elif pressed(imgui.Key.X):
|
||||
vis_options.flags[mujoco.mjtVisFlag.mjVIS_TEXTURE] ^= 1
|
||||
elif pressed(imgui.Key.J):
|
||||
vis_options.flags[mujoco.mjtVisFlag.mjVIS_JOINT] ^= 1
|
||||
elif not is_freecam_wasd and pressed(imgui.Key.Q):
|
||||
vis_options.flags[mujoco.mjtVisFlag.mjVIS_CAMERA] ^= 1
|
||||
elif pressed(imgui.Key.U):
|
||||
vis_options.flags[mujoco.mjtVisFlag.mjVIS_ACTUATOR] ^= 1
|
||||
elif pressed(imgui.Key.Comma):
|
||||
vis_options.flags[mujoco.mjtVisFlag.mjVIS_ACTIVATION] ^= 1
|
||||
elif pressed(imgui.Key.Z):
|
||||
vis_options.flags[mujoco.mjtVisFlag.mjVIS_LIGHT] ^= 1
|
||||
elif pressed(imgui.Key.V):
|
||||
vis_options.flags[mujoco.mjtVisFlag.mjVIS_TENDON] ^= 1
|
||||
elif pressed(imgui.Key.Y):
|
||||
vis_options.flags[mujoco.mjtVisFlag.mjVIS_RANGEFINDER] ^= 1
|
||||
elif not is_freecam_wasd and pressed(imgui.Key.E):
|
||||
vis_options.flags[mujoco.mjtVisFlag.mjVIS_CONSTRAINT] ^= 1
|
||||
elif pressed(imgui.Key.I):
|
||||
vis_options.flags[mujoco.mjtVisFlag.mjVIS_INERTIA] ^= 1
|
||||
elif pressed(imgui.Key.Apostrophe):
|
||||
vis_options.flags[mujoco.mjtVisFlag.mjVIS_SCLINERTIA] ^= 1
|
||||
elif pressed(imgui.Key.B):
|
||||
vis_options.flags[mujoco.mjtVisFlag.mjVIS_PERTFORCE] ^= 1
|
||||
elif pressed(imgui.Key.O):
|
||||
vis_options.flags[mujoco.mjtVisFlag.mjVIS_PERTOBJ] ^= 1
|
||||
elif pressed(imgui.Key.C):
|
||||
vis_options.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] ^= 1
|
||||
elif pressed(imgui.Key.N):
|
||||
vis_options.flags[mujoco.mjtVisFlag.mjVIS_ISLAND] ^= 1
|
||||
elif pressed(imgui.Key.F):
|
||||
vis_options.flags[mujoco.mjtVisFlag.mjVIS_CONTACTFORCE] ^= 1
|
||||
elif pressed(imgui.Key.P):
|
||||
vis_options.flags[mujoco.mjtVisFlag.mjVIS_CONTACTSPLIT] ^= 1
|
||||
elif pressed(imgui.Key.T):
|
||||
vis_options.flags[mujoco.mjtVisFlag.mjVIS_TRANSPARENT] ^= 1
|
||||
elif not is_freecam_wasd and pressed(imgui.Key.A):
|
||||
vis_options.flags[mujoco.mjtVisFlag.mjVIS_AUTOCONNECT] ^= 1
|
||||
elif pressed(imgui.Key.M):
|
||||
vis_options.flags[mujoco.mjtVisFlag.mjVIS_COM] ^= 1
|
||||
elif not is_freecam_wasd and pressed(imgui.Key.D):
|
||||
vis_options.flags[mujoco.mjtVisFlag.mjVIS_STATIC] ^= 1
|
||||
elif pressed(imgui.Key.Semicolon):
|
||||
vis_options.flags[mujoco.mjtVisFlag.mjVIS_SKIN] ^= 1
|
||||
elif pressed(imgui.Key.GraveAccent):
|
||||
vis_options.flags[mujoco.mjtVisFlag.mjVIS_BODYBVH] ^= 1
|
||||
elif pressed(imgui.Key.Backslash):
|
||||
vis_options.flags[mujoco.mjtVisFlag.mjVIS_MESHBVH] ^= 1
|
||||
|
||||
# Site group toggles (Shift + 0-5).
|
||||
elif pressed(int(imgui.Key.N0) | int(imgui.Key.Shift)):
|
||||
vis_options.sitegroup[0] ^= 1
|
||||
elif pressed(int(imgui.Key.N1) | int(imgui.Key.Shift)):
|
||||
vis_options.sitegroup[1] ^= 1
|
||||
elif pressed(int(imgui.Key.N2) | int(imgui.Key.Shift)):
|
||||
vis_options.sitegroup[2] ^= 1
|
||||
elif pressed(int(imgui.Key.N3) | int(imgui.Key.Shift)):
|
||||
vis_options.sitegroup[3] ^= 1
|
||||
elif pressed(int(imgui.Key.N4) | int(imgui.Key.Shift)):
|
||||
vis_options.sitegroup[4] ^= 1
|
||||
elif pressed(int(imgui.Key.N5) | int(imgui.Key.Shift)):
|
||||
vis_options.sitegroup[5] ^= 1
|
||||
|
||||
# Geom group toggles (0-5).
|
||||
elif pressed(imgui.Key.N0):
|
||||
vis_options.geomgroup[0] ^= 1
|
||||
elif pressed(imgui.Key.N1):
|
||||
vis_options.geomgroup[1] ^= 1
|
||||
elif pressed(imgui.Key.N2):
|
||||
vis_options.geomgroup[2] ^= 1
|
||||
elif pressed(imgui.Key.N3):
|
||||
vis_options.geomgroup[3] ^= 1
|
||||
elif pressed(imgui.Key.N4):
|
||||
vis_options.geomgroup[4] ^= 1
|
||||
elif pressed(imgui.Key.N5):
|
||||
vis_options.geomgroup[5] ^= 1
|
||||
|
||||
else:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def handle_step_control_keyboard_events(
|
||||
model: mujoco.MjModel,
|
||||
data: mujoco.MjData,
|
||||
step_control: sim.StepControl,
|
||||
ux_state: ux.UxState,
|
||||
) -> bool:
|
||||
"""Handles keyboard shortcuts for simulation stepping control.
|
||||
|
||||
Args:
|
||||
model: The MuJoCo model.
|
||||
data: The MuJoCo data.
|
||||
step_control: The simulation step control object.
|
||||
ux_state: The UX state object.
|
||||
|
||||
Returns:
|
||||
True if a key was handled, False otherwise.
|
||||
"""
|
||||
if imgui.GetIO().WantCaptureKeyboard:
|
||||
return False
|
||||
|
||||
pressed = imgui.IsKeyChordPressed
|
||||
|
||||
if pressed(int(imgui.Key.Ctrl) | int(imgui.Key.Space)):
|
||||
if step_control.get_pause_state() == sim.PauseState.VISCOUS_PAUSED:
|
||||
step_control.set_pause_state(sim.PauseState.UNPAUSED)
|
||||
else:
|
||||
step_control.set_pause_state(sim.PauseState.VISCOUS_PAUSED)
|
||||
return True
|
||||
elif pressed(imgui.Key.Space):
|
||||
pause = step_control.get_pause_state()
|
||||
if pause in (sim.PauseState.VISCOUS_PAUSED, sim.PauseState.UNPAUSED):
|
||||
step_control.set_pause_state(sim.PauseState.NORMAL_PAUSED)
|
||||
else:
|
||||
step_control.set_pause_state(sim.PauseState.UNPAUSED)
|
||||
return True
|
||||
elif pressed(imgui.Key.Backspace):
|
||||
mujoco.mj_resetData(model, data)
|
||||
mujoco.mj_forward(model, data)
|
||||
return True
|
||||
elif pressed(imgui.Key.Minus):
|
||||
ux.set_speed_index(step_control, ux_state, ux_state.speed_index + 1)
|
||||
return True
|
||||
elif pressed(imgui.Key.Equal):
|
||||
ux.set_speed_index(step_control, ux_state, ux_state.speed_index - 1)
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def handle_camera_select_keyboard_events(
|
||||
model: mujoco.MjModel,
|
||||
camera: mujoco.MjvCamera,
|
||||
ux_state: ux.UxState,
|
||||
) -> bool:
|
||||
"""Handles keyboard shortcuts for camera selection.
|
||||
|
||||
Args:
|
||||
model: The MuJoCo model.
|
||||
camera: The MuJoCo camera object.
|
||||
ux_state: The UX state object.
|
||||
|
||||
Returns:
|
||||
True if a key was handled, False otherwise.
|
||||
"""
|
||||
if imgui.GetIO().WantCaptureKeyboard:
|
||||
return False
|
||||
|
||||
pressed = imgui.IsKeyChordPressed
|
||||
|
||||
if pressed(imgui.Key.Escape):
|
||||
ux_state.camera_index = ux.set_camera(model, camera, ux.TUMBLE_CAMERA_IDX)
|
||||
return True
|
||||
elif pressed(imgui.Key.LeftBracket):
|
||||
ux_state.camera_index = ux.set_camera(
|
||||
model, camera, ux_state.camera_index - 1
|
||||
)
|
||||
return True
|
||||
elif pressed(imgui.Key.RightBracket):
|
||||
ux_state.camera_index = ux.set_camera(
|
||||
model, camera, ux_state.camera_index + 1
|
||||
)
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def handle_freecam_wasd_keyboard_events(
|
||||
model: mujoco.MjModel,
|
||||
data: mujoco.MjData,
|
||||
camera: mujoco.MjvCamera,
|
||||
cam_speed: float,
|
||||
) -> tuple[bool, float]:
|
||||
"""Handles keyboard shortcuts for free camera movement.
|
||||
|
||||
Args:
|
||||
model: The MuJoCo model.
|
||||
data: The MuJoCo data.
|
||||
camera: The MuJoCo camera object.
|
||||
cam_speed: The current camera speed.
|
||||
|
||||
Returns:
|
||||
A tuple of (handled, updated_cam_speed).
|
||||
"""
|
||||
if imgui.GetIO().WantCaptureKeyboard:
|
||||
return False, cam_speed
|
||||
|
||||
if (
|
||||
imgui.IsKeyDown(imgui.Key.W)
|
||||
or imgui.IsKeyDown(imgui.Key.S)
|
||||
or imgui.IsKeyDown(imgui.Key.A)
|
||||
or imgui.IsKeyDown(imgui.Key.D)
|
||||
or imgui.IsKeyDown(imgui.Key.Q)
|
||||
or imgui.IsKeyDown(imgui.Key.E)
|
||||
):
|
||||
moved = False
|
||||
|
||||
# Move (dolly) forward/backward using W and S keys.
|
||||
if imgui.IsKeyDown(imgui.Key.W):
|
||||
ux.MoveCamera(
|
||||
model,
|
||||
data,
|
||||
camera,
|
||||
ux.CameraMotion.TRUCK_DOLLY,
|
||||
0,
|
||||
cam_speed,
|
||||
)
|
||||
moved = True
|
||||
elif imgui.IsKeyDown(imgui.Key.S):
|
||||
ux.MoveCamera(
|
||||
model,
|
||||
data,
|
||||
camera,
|
||||
ux.CameraMotion.TRUCK_DOLLY,
|
||||
0,
|
||||
-cam_speed,
|
||||
)
|
||||
moved = True
|
||||
|
||||
# Strafe (truck) left/right using A and D keys.
|
||||
if imgui.IsKeyDown(imgui.Key.A):
|
||||
ux.MoveCamera(
|
||||
model,
|
||||
data,
|
||||
camera,
|
||||
ux.CameraMotion.TRUCK_DOLLY,
|
||||
-cam_speed,
|
||||
0,
|
||||
)
|
||||
moved = True
|
||||
elif imgui.IsKeyDown(imgui.Key.D):
|
||||
ux.MoveCamera(
|
||||
model,
|
||||
data,
|
||||
camera,
|
||||
ux.CameraMotion.TRUCK_DOLLY,
|
||||
cam_speed,
|
||||
0,
|
||||
)
|
||||
moved = True
|
||||
|
||||
# Move (pedestal) up/down using Q and E keys.
|
||||
if imgui.IsKeyDown(imgui.Key.Q):
|
||||
ux.MoveCamera(
|
||||
model,
|
||||
data,
|
||||
camera,
|
||||
ux.CameraMotion.TRUCK_PEDESTAL,
|
||||
0,
|
||||
cam_speed,
|
||||
)
|
||||
moved = True
|
||||
elif imgui.IsKeyDown(imgui.Key.E):
|
||||
ux.MoveCamera(
|
||||
model,
|
||||
data,
|
||||
camera,
|
||||
ux.CameraMotion.TRUCK_PEDESTAL,
|
||||
0,
|
||||
-cam_speed,
|
||||
)
|
||||
moved = True
|
||||
|
||||
if moved:
|
||||
cam_speed += 0.001
|
||||
max_speed = 0.1 if imgui.GetIO().KeyShift else 0.01
|
||||
if cam_speed > max_speed:
|
||||
cam_speed = max_speed
|
||||
else:
|
||||
cam_speed = 0.001
|
||||
|
||||
return True, cam_speed
|
||||
|
||||
return False, cam_speed
|
||||
|
||||
|
||||
def handle_keyboard_events(
|
||||
model: mujoco.MjModel,
|
||||
data: mujoco.MjData,
|
||||
camera: mujoco.MjvCamera,
|
||||
vis_options: mujoco.MjvOption,
|
||||
step_control: sim.StepControl,
|
||||
ux_state: ux.UxState,
|
||||
cam_speed: float,
|
||||
) -> tuple[bool, float]:
|
||||
"""Handle keyboard events according to Studio's bindings.
|
||||
|
||||
Args:
|
||||
model: The MuJoCo model.
|
||||
data: The MuJoCo data.
|
||||
camera: The MuJoCo camera object.
|
||||
vis_options: The MuJoCo visualization options.
|
||||
step_control: The simulation step control object.
|
||||
ux_state: The UX state object.
|
||||
cam_speed: The current camera speed.
|
||||
|
||||
Returns:
|
||||
A tuple of (handled, updated_cam_speed).
|
||||
"""
|
||||
if imgui.GetIO().WantCaptureKeyboard:
|
||||
return False, cam_speed
|
||||
|
||||
is_freecam_wasd = ux_state.camera_index == ux.FREE_CAMERA_IDX
|
||||
if handle_step_control_keyboard_events(model, data, step_control, ux_state):
|
||||
return True, cam_speed
|
||||
|
||||
if handle_camera_select_keyboard_events(model, camera, ux_state):
|
||||
return True, cam_speed
|
||||
|
||||
if handle_vis_options_keyboard_events(vis_options, is_freecam_wasd):
|
||||
return True, cam_speed
|
||||
|
||||
if is_freecam_wasd:
|
||||
handled, cam_speed = handle_freecam_wasd_keyboard_events(
|
||||
model, data, camera, cam_speed
|
||||
)
|
||||
if handled:
|
||||
return True, cam_speed
|
||||
|
||||
return False, cam_speed
|
||||
|
||||
|
||||
def handle_camera_tracking_mouse_events(
|
||||
model: mujoco.MjModel,
|
||||
data: mujoco.MjData,
|
||||
camera: mujoco.MjvCamera,
|
||||
vis_options: mujoco.MjvOption,
|
||||
ux_state: ux.UxState,
|
||||
) -> None:
|
||||
"""Handles mouse events for camera tracking."""
|
||||
io = imgui.GetIO()
|
||||
if imgui.GetIO().WantCaptureMouse:
|
||||
return
|
||||
|
||||
if io.DisplaySize.x <= 0 or io.DisplaySize.y <= 0:
|
||||
return
|
||||
|
||||
mouse_x = io.MousePos.x / io.DisplaySize.x
|
||||
mouse_y = io.MousePos.y / io.DisplaySize.y
|
||||
aspect_ratio = io.DisplaySize.x / io.DisplaySize.y
|
||||
|
||||
# Right double click.
|
||||
if imgui.IsMouseDoubleClicked(imgui.MouseButton.Right):
|
||||
picked = ux.Pick(
|
||||
model,
|
||||
data,
|
||||
camera,
|
||||
mouse_x,
|
||||
mouse_y,
|
||||
aspect_ratio,
|
||||
vis_options,
|
||||
)
|
||||
if picked.body > 0 and io.KeyCtrl:
|
||||
# Switch camera to tracking mode and track the selected body.
|
||||
camera.type = int(mujoco.mjtCamera.mjCAMERA_TRACKING)
|
||||
camera.trackbodyid = picked.body
|
||||
camera.fixedcamid = -1
|
||||
ux_state.camera_index = ux.TRACKING_CAMERA_IDX
|
||||
|
||||
|
||||
def handle_mouse_events(
|
||||
model: mujoco.MjModel,
|
||||
data: mujoco.MjData,
|
||||
camera: mujoco.MjvCamera,
|
||||
vis_options: mujoco.MjvOption,
|
||||
perturb: mujoco.MjvPerturb,
|
||||
ux_state: ux.UxState,
|
||||
) -> None:
|
||||
"""Handles mouse events."""
|
||||
io = imgui.GetIO()
|
||||
if io.WantCaptureMouse:
|
||||
return
|
||||
|
||||
if io.DisplaySize.x <= 0 or io.DisplaySize.y <= 0:
|
||||
return
|
||||
|
||||
mouse_x = io.MousePos.x / io.DisplaySize.x
|
||||
mouse_y = io.MousePos.y / io.DisplaySize.y
|
||||
mouse_dx = io.MouseDelta.x / io.DisplaySize.x
|
||||
mouse_dy = io.MouseDelta.y / io.DisplaySize.y
|
||||
mouse_scroll = io.MouseWheel / 50.0
|
||||
|
||||
is_mouse_moving = mouse_dx != 0.0 or mouse_dy != 0.0
|
||||
is_any_mouse_down = (
|
||||
imgui.IsMouseDown(imgui.MouseButton.Left)
|
||||
or imgui.IsMouseDown(imgui.MouseButton.Right)
|
||||
or imgui.IsMouseDown(imgui.MouseButton.Middle)
|
||||
)
|
||||
is_mouse_dragging = is_mouse_moving and is_any_mouse_down
|
||||
|
||||
# If no mouse buttons are down, end any active perturbations.
|
||||
if not is_any_mouse_down:
|
||||
perturb.active = 0
|
||||
|
||||
# Handle perturbation mouse actions.
|
||||
if is_mouse_dragging and io.KeyCtrl:
|
||||
if perturb.select > 0:
|
||||
action = int(mujoco.mjtMouse.mjMOUSE_NONE)
|
||||
if imgui.IsMouseDown(imgui.MouseButton.Left):
|
||||
if io.KeyAlt:
|
||||
action = int(
|
||||
mujoco.mjtMouse.mjMOUSE_MOVE_H
|
||||
if io.KeyShift
|
||||
else mujoco.mjtMouse.mjMOUSE_MOVE_V
|
||||
)
|
||||
else:
|
||||
action = int(
|
||||
mujoco.mjtMouse.mjMOUSE_ROTATE_H
|
||||
if io.KeyShift
|
||||
else mujoco.mjtMouse.mjMOUSE_ROTATE_V
|
||||
)
|
||||
elif imgui.IsMouseDown(imgui.MouseButton.Right):
|
||||
action = int(
|
||||
mujoco.mjtMouse.mjMOUSE_MOVE_H
|
||||
if io.KeyShift
|
||||
else mujoco.mjtMouse.mjMOUSE_MOVE_V
|
||||
)
|
||||
elif imgui.IsMouseDown(imgui.MouseButton.Middle):
|
||||
action = int(mujoco.mjtMouse.mjMOUSE_ZOOM)
|
||||
|
||||
active = int(
|
||||
mujoco.mjtPertBit.mjPERT_TRANSLATE
|
||||
if action
|
||||
in (
|
||||
int(mujoco.mjtMouse.mjMOUSE_MOVE_V),
|
||||
int(mujoco.mjtMouse.mjMOUSE_MOVE_H),
|
||||
)
|
||||
else mujoco.mjtPertBit.mjPERT_ROTATE
|
||||
)
|
||||
if active != perturb.active:
|
||||
ux.InitPerturb(model, data, camera, perturb, active)
|
||||
ux.MovePerturb(
|
||||
model,
|
||||
data,
|
||||
camera,
|
||||
perturb,
|
||||
action,
|
||||
mouse_dx,
|
||||
mouse_dy,
|
||||
)
|
||||
elif is_mouse_dragging:
|
||||
if ux_state.camera_index == ux.FREE_CAMERA_IDX:
|
||||
if imgui.IsMouseDown(imgui.MouseButton.Left):
|
||||
ux.MoveCamera(
|
||||
model,
|
||||
data,
|
||||
camera,
|
||||
ux.CameraMotion.PAN_TILT,
|
||||
mouse_dx,
|
||||
mouse_dy,
|
||||
)
|
||||
else:
|
||||
if imgui.IsMouseDown(imgui.MouseButton.Left):
|
||||
ux.MoveCamera(
|
||||
model,
|
||||
data,
|
||||
camera,
|
||||
ux.CameraMotion.ORBIT,
|
||||
mouse_dx,
|
||||
mouse_dy,
|
||||
)
|
||||
elif imgui.IsMouseDown(imgui.MouseButton.Middle):
|
||||
ux.MoveCamera(
|
||||
model,
|
||||
data,
|
||||
camera,
|
||||
ux.CameraMotion.ZOOM,
|
||||
mouse_dx,
|
||||
mouse_dy,
|
||||
)
|
||||
|
||||
# Right mouse movement is relative to the horizontal and vertical planes.
|
||||
if imgui.IsMouseDown(imgui.MouseButton.Right) and io.KeyShift:
|
||||
ux.MoveCamera(
|
||||
model,
|
||||
data,
|
||||
camera,
|
||||
ux.CameraMotion.PLANAR_MOVE_H,
|
||||
mouse_dx,
|
||||
mouse_dy,
|
||||
)
|
||||
elif imgui.IsMouseDown(imgui.MouseButton.Right):
|
||||
ux.MoveCamera(
|
||||
model,
|
||||
data,
|
||||
camera,
|
||||
ux.CameraMotion.PLANAR_MOVE_V,
|
||||
mouse_dx,
|
||||
mouse_dy,
|
||||
)
|
||||
|
||||
# Mouse scroll.
|
||||
if mouse_scroll != 0.0 and ux_state.camera_index != ux.FREE_CAMERA_IDX:
|
||||
ux.MoveCamera(
|
||||
model,
|
||||
data,
|
||||
camera,
|
||||
ux.CameraMotion.ZOOM,
|
||||
0,
|
||||
-mouse_scroll,
|
||||
)
|
||||
|
||||
aspect_ratio = io.DisplaySize.x / io.DisplaySize.y
|
||||
|
||||
# Left double click.
|
||||
if imgui.IsMouseDoubleClicked(imgui.MouseButton.Left):
|
||||
picked = ux.Pick(
|
||||
model,
|
||||
data,
|
||||
camera,
|
||||
mouse_x,
|
||||
mouse_y,
|
||||
aspect_ratio,
|
||||
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.
|
||||
tmp = np.array(picked.point, dtype=np.float64) - data.xpos[picked.body]
|
||||
xmat = np.array(data.xmat[picked.body], dtype=np.float64).reshape(3, 3)
|
||||
perturb.localpos = xmat.T @ tmp
|
||||
else:
|
||||
perturb.select = 0
|
||||
perturb.flexselect = -1
|
||||
perturb.skinselect = -1
|
||||
|
||||
handle_camera_tracking_mouse_events(
|
||||
model, data, camera, vis_options, ux_state
|
||||
)
|
||||
@@ -0,0 +1,398 @@
|
||||
// Copyright 2026 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
|
||||
//
|
||||
// https://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.
|
||||
|
||||
// Python bindings for MuJoCo platform UX components.
|
||||
|
||||
#include <algorithm>
|
||||
#include <array>
|
||||
#include <string>
|
||||
#include <tuple>
|
||||
#include <vector>
|
||||
|
||||
#include <imgui.h>
|
||||
#include <mujoco/mujoco.h>
|
||||
#include "third_party/mujoco/src/experimental/platform/helpers.h"
|
||||
#include "third_party/mujoco/src/experimental/platform/sim/step_control.h"
|
||||
#include "third_party/mujoco/src/experimental/platform/ux/gui.h"
|
||||
#include "third_party/mujoco/src/experimental/platform/ux/interaction.h"
|
||||
#include "structs.h"
|
||||
#include <pybind11/pybind11.h>
|
||||
#include <pybind11/stl.h>
|
||||
|
||||
namespace py = pybind11;
|
||||
|
||||
struct UxState {
|
||||
// Read/edited by step_control_gui
|
||||
int speed_index = 0;
|
||||
|
||||
// Read/edited by state_gui
|
||||
std::vector<mjtNum> state;
|
||||
int state_sig = 0;
|
||||
|
||||
// Read/edited by watch_gui
|
||||
char watch_field_name[256] = {0};
|
||||
int watch_field_index = 0;
|
||||
|
||||
// Read/edited by noise_gui
|
||||
float noise_scale = 0.0f;
|
||||
float noise_rate = 0.0f;
|
||||
|
||||
// Read/edited by camera_selection_gui
|
||||
int camera_index = mujoco::platform::kTumbleCameraIdx;
|
||||
};
|
||||
|
||||
struct RenderFlags {
|
||||
std::array<uint8_t, mjNRNDFLAG> flags = {0};
|
||||
};
|
||||
|
||||
PYBIND11_MODULE(ux, m) {
|
||||
py::class_<RenderFlags>(m, "RenderFlags")
|
||||
.def(py::init<>())
|
||||
.def_readwrite("flags", &RenderFlags::flags);
|
||||
|
||||
m.doc() = "MuJoCo platform UX components.";
|
||||
|
||||
py::enum_<mujoco::platform::GuiTheme>(m, "GuiTheme")
|
||||
.value("LIGHT", mujoco::platform::GuiTheme::kLight)
|
||||
.value("DARK", mujoco::platform::GuiTheme::kDark)
|
||||
.value("CLASSIC", mujoco::platform::GuiTheme::kClassic);
|
||||
|
||||
py::class_<UxState>(m, "UxState")
|
||||
.def(py::init<>())
|
||||
.def_readwrite("speed_index", &UxState::speed_index)
|
||||
.def_readwrite("state", &UxState::state)
|
||||
.def_readwrite("state_sig", &UxState::state_sig)
|
||||
.def_readwrite("watch_field_index", &UxState::watch_field_index)
|
||||
.def_readwrite("noise_scale", &UxState::noise_scale)
|
||||
.def_readwrite("noise_rate", &UxState::noise_rate)
|
||||
.def_readwrite("camera_index", &UxState::camera_index)
|
||||
.def_property(
|
||||
"watch_field_name",
|
||||
[](const UxState& self) {
|
||||
return std::string(self.watch_field_name);
|
||||
},
|
||||
[](UxState& self, const std::string& val) {
|
||||
std::snprintf(self.watch_field_name, sizeof(self.watch_field_name),
|
||||
"%s", val.c_str());
|
||||
});
|
||||
|
||||
m.def(
|
||||
"setup_theme",
|
||||
[](mujoco::platform::GuiTheme theme) {
|
||||
mujoco::platform::SetupTheme(theme);
|
||||
},
|
||||
py::arg("theme"), "Set up Dear ImGui visual theme.");
|
||||
|
||||
m.def(
|
||||
"configure_docking_layout",
|
||||
[]() {
|
||||
ImVec4 r = mujoco::platform::ConfigureDockingLayout();
|
||||
return std::make_tuple(r.x, r.y, r.z, r.w);
|
||||
},
|
||||
"Configure the docking layout with Options (left) and Inspector (right) "
|
||||
"panes. Returns (x, y, w, h) of the central workspace area.");
|
||||
|
||||
m.def(
|
||||
"step_control_gui",
|
||||
[](const mujoco::python::MjModelWrapper& model,
|
||||
mujoco::platform::StepControl* step_control, UxState& ux_state) {
|
||||
mujoco::platform::StepControlGui(model.get(), step_control,
|
||||
ux_state.speed_index);
|
||||
},
|
||||
py::arg("model"), py::arg("step_control"), py::arg("ux_state"),
|
||||
"Render the simulation stepping control GUI. Modifies "
|
||||
"ux_state.speed_index.");
|
||||
|
||||
m.def(
|
||||
"theme_select_gui",
|
||||
[](mujoco::platform::GuiTheme theme) {
|
||||
bool changed = mujoco::platform::ThemeSelectGui(&theme);
|
||||
return std::make_tuple(changed, theme);
|
||||
},
|
||||
py::arg("theme"),
|
||||
"Render the GUI theme selector. Returns (changed, theme).");
|
||||
|
||||
m.def(
|
||||
"label_selection_gui",
|
||||
[](mujoco::python::MjvOptionWrapper& vis_options) {
|
||||
return mujoco::platform::LabelSelectionGui(vis_options.get());
|
||||
},
|
||||
py::arg("vis_options"), "Render the visualization label selection GUI.");
|
||||
|
||||
m.def(
|
||||
"frame_selection_gui",
|
||||
[](mujoco::python::MjvOptionWrapper& vis_options) {
|
||||
return mujoco::platform::FrameSelectionGui(vis_options.get());
|
||||
},
|
||||
py::arg("vis_options"), "Render the visualization frame selection GUI.");
|
||||
|
||||
m.def(
|
||||
"camera_selection_gui",
|
||||
[](const mujoco::python::MjModelWrapper& model,
|
||||
mujoco::python::MjDataWrapper& data,
|
||||
mujoco::python::MjvCameraWrapper& camera, UxState& ux_state) {
|
||||
bool changed = mujoco::platform::CameraSelectionGui(
|
||||
model.get(), data.get(), *camera.get(), ux_state.camera_index);
|
||||
return changed;
|
||||
},
|
||||
py::arg("model"), py::arg("data"), py::arg("camera"), py::arg("ux_state"),
|
||||
"Render the camera selection GUI. Modifies ux_state.camera_index. "
|
||||
"Returns true if camera changed.");
|
||||
|
||||
m.def(
|
||||
"set_camera",
|
||||
[](const mujoco::python::MjModelWrapper& model,
|
||||
mujoco::python::MjvCameraWrapper& camera, int request_idx) {
|
||||
return mujoco::platform::SetCamera(model.get(), camera.get(), request_idx);
|
||||
},
|
||||
py::arg("model"), py::arg("camera"), py::arg("request_idx"),
|
||||
"Set the camera index and update the camera object.");
|
||||
|
||||
m.def(
|
||||
"set_speed_index",
|
||||
[](mujoco::platform::StepControl* step_control, UxState& ux_state, int idx) {
|
||||
mujoco::platform::SetSpeedIndex(step_control, ux_state.speed_index, idx);
|
||||
},
|
||||
py::arg("step_control"), py::arg("ux_state"), py::arg("idx"),
|
||||
"Set the simulation speed index.");
|
||||
|
||||
m.def(
|
||||
"physics_gui",
|
||||
[](mujoco::python::MjModelWrapper& model, float min_width) {
|
||||
mujoco::platform::PhysicsGui(model.get(), min_width);
|
||||
},
|
||||
py::arg("model"), py::arg("min_width") = 150.0f,
|
||||
"Render the physics settings UI.");
|
||||
|
||||
m.def(
|
||||
"rendering_gui",
|
||||
[](const mujoco::python::MjModelWrapper& model,
|
||||
mujoco::python::MjvOptionWrapper& vis_options,
|
||||
RenderFlags& render_flags) {
|
||||
mjtByte flags[mjNRNDFLAG] = {0};
|
||||
for (int i = 0; i < mjNRNDFLAG; ++i) {
|
||||
flags[i] = render_flags.flags[i];
|
||||
}
|
||||
mujoco::platform::RenderingGui(model.get(), vis_options.get(), flags,
|
||||
150.0f);
|
||||
for (int i = 0; i < mjNRNDFLAG; ++i) {
|
||||
render_flags.flags[i] = flags[i];
|
||||
}
|
||||
},
|
||||
py::arg("model"), py::arg("vis_options"), py::arg("render_flags"),
|
||||
"Render the rendering settings UI. Modifies render_flags in place.");
|
||||
|
||||
m.def(
|
||||
"groups_gui",
|
||||
[](const mujoco::python::MjModelWrapper& model,
|
||||
mujoco::python::MjvOptionWrapper& vis_options, float min_width) {
|
||||
mujoco::platform::GroupsGui(model.get(), vis_options.get(), min_width);
|
||||
},
|
||||
py::arg("model"), py::arg("vis_options"), py::arg("min_width") = 150.0f,
|
||||
"Render the visibility groups UI.");
|
||||
|
||||
m.def(
|
||||
"visualization_gui",
|
||||
[](mujoco::python::MjModelWrapper& model,
|
||||
mujoco::python::MjvOptionWrapper& vis_options,
|
||||
mujoco::python::MjvCameraWrapper& camera, float min_width) {
|
||||
mujoco::platform::VisualizationGui(model.get(), vis_options.get(),
|
||||
camera.get(), min_width);
|
||||
},
|
||||
py::arg("model"), py::arg("vis_options"), py::arg("camera"),
|
||||
py::arg("min_width") = 150.0f, "Render the visualization settings UI.");
|
||||
|
||||
m.def(
|
||||
"controls_gui",
|
||||
[](const mujoco::python::MjModelWrapper& model,
|
||||
mujoco::python::MjDataWrapper& data,
|
||||
mujoco::python::MjvOptionWrapper& vis_options) {
|
||||
mujoco::platform::ControlsGui(model.get(), data.get(),
|
||||
vis_options.get());
|
||||
},
|
||||
py::arg("model"), py::arg("data"), py::arg("vis_options"),
|
||||
"Render the actuator controls UI.");
|
||||
|
||||
m.def(
|
||||
"joints_gui",
|
||||
[](const mujoco::python::MjModelWrapper& model,
|
||||
mujoco::python::MjDataWrapper& data,
|
||||
mujoco::python::MjvOptionWrapper& vis_options) {
|
||||
mujoco::platform::JointsGui(model.get(), data.get(), vis_options.get());
|
||||
},
|
||||
py::arg("model"), py::arg("data"), py::arg("vis_options"),
|
||||
"Render the joints UI.");
|
||||
|
||||
m.def(
|
||||
"sensor_gui",
|
||||
[](const mujoco::python::MjModelWrapper& model,
|
||||
mujoco::python::MjDataWrapper& data) {
|
||||
mujoco::platform::SensorGui(model.get(), data.get());
|
||||
},
|
||||
py::arg("model"), py::arg("data"), "Render the sensor data plot.");
|
||||
|
||||
m.def(
|
||||
"state_gui",
|
||||
[](const mujoco::python::MjModelWrapper& model,
|
||||
mujoco::python::MjDataWrapper& data, UxState& ux_state,
|
||||
float min_width) {
|
||||
mujoco::platform::StateGui(model.get(), data.get(), ux_state.state,
|
||||
ux_state.state_sig, min_width);
|
||||
},
|
||||
py::arg("model"), py::arg("data"), py::arg("ux_state"),
|
||||
py::arg("min_width") = 150.0f,
|
||||
"Render the state UI. Modifies ux_state.state and ux_state.state_sig.");
|
||||
|
||||
m.def(
|
||||
"watch_gui",
|
||||
[](const mujoco::python::MjModelWrapper& model,
|
||||
mujoco::python::MjDataWrapper& data, UxState& ux_state) {
|
||||
mujoco::platform::WatchGui(
|
||||
model.get(), data.get(), ux_state.watch_field_name,
|
||||
sizeof(ux_state.watch_field_name), ux_state.watch_field_index);
|
||||
},
|
||||
py::arg("model"), py::arg("data"), py::arg("ux_state"),
|
||||
"Render the watch UI. Modifies ux_state.watch_field_name and "
|
||||
"ux_state.watch_field_index.");
|
||||
|
||||
m.def(
|
||||
"noise_gui",
|
||||
[](const mujoco::python::MjModelWrapper& model,
|
||||
mujoco::python::MjDataWrapper& data, UxState& ux_state) {
|
||||
mujoco::platform::NoiseGui(model.get(), data.get(),
|
||||
ux_state.noise_scale, ux_state.noise_rate);
|
||||
},
|
||||
py::arg("model"), py::arg("data"), py::arg("ux_state"),
|
||||
"Render the noise UI. Modifies ux_state.noise_scale and "
|
||||
"ux_state.noise_rate.");
|
||||
|
||||
m.def(
|
||||
"convergence_gui",
|
||||
[](const mujoco::python::MjModelWrapper& model,
|
||||
mujoco::python::MjDataWrapper& data) {
|
||||
mujoco::platform::ConvergenceGui(model.get(), data.get());
|
||||
},
|
||||
py::arg("model"), py::arg("data"),
|
||||
"Render the solver convergence chart.");
|
||||
|
||||
m.def(
|
||||
"counts_gui",
|
||||
[](const mujoco::python::MjModelWrapper& model,
|
||||
mujoco::python::MjDataWrapper& data) {
|
||||
mujoco::platform::CountsGui(model.get(), data.get());
|
||||
},
|
||||
py::arg("model"), py::arg("data"), "Render the solver counts chart.");
|
||||
|
||||
m.def(
|
||||
"stats_gui",
|
||||
[](const mujoco::python::MjModelWrapper& model,
|
||||
mujoco::python::MjDataWrapper& data, bool paused, float fps) {
|
||||
mujoco::platform::StatsGui(model.get(), data.get(), paused, fps);
|
||||
},
|
||||
py::arg("model"), py::arg("data"), py::arg("paused"), py::arg("fps"),
|
||||
"Render the simulation statistics UI.");
|
||||
|
||||
m.attr("FREE_CAMERA_IDX") = mujoco::platform::kFreeCameraIdx;
|
||||
m.attr("TUMBLE_CAMERA_IDX") = mujoco::platform::kTumbleCameraIdx;
|
||||
m.attr("TRACKING_CAMERA_IDX") = mujoco::platform::kTrackingCameraIdx;
|
||||
|
||||
py::enum_<mujoco::platform::CameraMotion>(m, "CameraMotion")
|
||||
.value("ZOOM", mujoco::platform::CameraMotion::ZOOM)
|
||||
.value("ORBIT", mujoco::platform::CameraMotion::ORBIT)
|
||||
.value("TRUCK_PEDESTAL", mujoco::platform::CameraMotion::TRUCK_PEDESTAL)
|
||||
.value("TRUCK_DOLLY", mujoco::platform::CameraMotion::TRUCK_DOLLY)
|
||||
.value("PAN_TILT", mujoco::platform::CameraMotion::PAN_TILT)
|
||||
.value("PLANAR_MOVE_H", mujoco::platform::CameraMotion::PLANAR_MOVE_H)
|
||||
.value("PLANAR_MOVE_V", mujoco::platform::CameraMotion::PLANAR_MOVE_V)
|
||||
.export_values();
|
||||
|
||||
m.def(
|
||||
"MoveCamera",
|
||||
[](const mujoco::python::MjModelWrapper& model,
|
||||
const mujoco::python::MjDataWrapper& data,
|
||||
mujoco::python::MjvCameraWrapper& cam,
|
||||
mujoco::platform::CameraMotion motion, mjtNum dx, mjtNum dy) {
|
||||
mujoco::platform::MoveCamera(model.get(), data.get(), cam.get(), motion,
|
||||
dx, dy);
|
||||
},
|
||||
py::arg("model"), py::arg("data"), py::arg("cam"), py::arg("motion"),
|
||||
py::arg("dx"), py::arg("dy"), "Moves the given camera.");
|
||||
|
||||
m.def(
|
||||
"InitPerturb",
|
||||
[](const mujoco::python::MjModelWrapper& model,
|
||||
const mujoco::python::MjDataWrapper& data,
|
||||
const mujoco::python::MjvCameraWrapper& cam,
|
||||
mujoco::python::MjvPerturbWrapper& pert, int active) {
|
||||
mujoco::platform::InitPerturb(model.get(), data.get(), cam.get(),
|
||||
pert.get(),
|
||||
static_cast<mjtPertBit>(active));
|
||||
},
|
||||
py::arg("model"), py::arg("data"), py::arg("cam"), py::arg("pert"),
|
||||
py::arg("active"), "Initializes mouse perturbation.");
|
||||
|
||||
m.def(
|
||||
"MovePerturb",
|
||||
[](const mujoco::python::MjModelWrapper& model,
|
||||
const mujoco::python::MjDataWrapper& data,
|
||||
const mujoco::python::MjvCameraWrapper& cam,
|
||||
mujoco::python::MjvPerturbWrapper& pert, int action, mjtNum reldx,
|
||||
mjtNum reldy) {
|
||||
mujoco::platform::MovePerturb(model.get(), data.get(), cam.get(),
|
||||
pert.get(), static_cast<mjtMouse>(action),
|
||||
reldx, reldy);
|
||||
},
|
||||
py::arg("model"), py::arg("data"), py::arg("cam"), py::arg("pert"),
|
||||
py::arg("action"), py::arg("reldx"), py::arg("reldy"),
|
||||
"Moves mouse perturbation.");
|
||||
|
||||
py::class_<mujoco::platform::PickResult>(m, "PickResult")
|
||||
.def_readwrite("dist", &mujoco::platform::PickResult::dist)
|
||||
.def_readwrite("body", &mujoco::platform::PickResult::body)
|
||||
.def_readwrite("geom", &mujoco::platform::PickResult::geom)
|
||||
.def_readwrite("flex", &mujoco::platform::PickResult::flex)
|
||||
.def_readwrite("skin", &mujoco::platform::PickResult::skin)
|
||||
.def_property(
|
||||
"point",
|
||||
[](const mujoco::platform::PickResult& res) {
|
||||
return py::make_tuple(res.point[0], res.point[1], res.point[2]);
|
||||
},
|
||||
[](mujoco::platform::PickResult& res, const py::tuple& t) {
|
||||
res.point[0] = t[0].cast<mjtNum>();
|
||||
res.point[1] = t[1].cast<mjtNum>();
|
||||
res.point[2] = t[2].cast<mjtNum>();
|
||||
});
|
||||
|
||||
m.def(
|
||||
"Pick",
|
||||
[](const mujoco::python::MjModelWrapper& model,
|
||||
const mujoco::python::MjDataWrapper& data,
|
||||
const mujoco::python::MjvCameraWrapper& cam, float x, float y,
|
||||
float aspect_ratio, const mujoco::python::MjvOptionWrapper& opt) {
|
||||
return mujoco::platform::Pick(model.get(), data.get(), cam.get(), x, y,
|
||||
aspect_ratio, opt.get());
|
||||
},
|
||||
py::arg("model"), py::arg("data"), py::arg("cam"), py::arg("x"),
|
||||
py::arg("y"), py::arg("aspect_ratio"), py::arg("opt"),
|
||||
"Picks object under cursor.");
|
||||
|
||||
m.def(
|
||||
"camera_to_string",
|
||||
[](const mujoco::python::MjDataWrapper& data,
|
||||
const mujoco::python::MjvCameraWrapper& camera) {
|
||||
return mujoco::platform::CameraToString(data.get(), camera.get());
|
||||
},
|
||||
py::arg("data"), py::arg("camera"),
|
||||
"Returns an XML string representation of the camera.");
|
||||
}
|
||||
@@ -0,0 +1,43 @@
|
||||
# Copyright 2026 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
|
||||
#
|
||||
# https://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.
|
||||
"""Structural protocol defining the common viewer interface.
|
||||
|
||||
StudioApp uses the protocol for convenience methods that accept any viewer.
|
||||
"""
|
||||
|
||||
from typing import Protocol
|
||||
|
||||
import mujoco
|
||||
from mujoco.experimental.studio import ux
|
||||
|
||||
|
||||
class Viewer(Protocol):
|
||||
"""Structural interface for any viewer."""
|
||||
|
||||
camera: mujoco.MjvCamera
|
||||
perturb: mujoco.MjvPerturb
|
||||
vis_options: mujoco.MjvOption
|
||||
render_flags: ux.RenderFlags
|
||||
|
||||
def is_running(self) -> bool:
|
||||
...
|
||||
|
||||
def sync(self, model: mujoco.MjModel, data: mujoco.MjData) -> None:
|
||||
...
|
||||
|
||||
def stop(self) -> None:
|
||||
...
|
||||
|
||||
def get_drop_file(self) -> str:
|
||||
...
|
||||
Reference in New Issue
Block a user