chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Move the Simulation panel into the platform toolkit.

Browse Source 
PiperOrigin-RevId: 965914672
Change-Id: Ib349b77055c57453780818dd4d8f48899b7024d5
This commit is contained in:
Yuval Tassa
2026-08-17 05:39:43 -07:00
committed by Copybara-Service
parent 0fa066237a
commit c232938233
7 changed files with 729 additions and 518 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/experimental/studio/sim.cc
+16
View File
@@ -17,6 +17,7 @@
#include <tuple>
#include <mujoco/mujoco.h>
#include <mujoco/experimental/platform/sim/sim_history.h>
#include <mujoco/experimental/platform/sim/step_control.h>
#include "structs.h"
#include <pybind11/pybind11.h>
@@ -82,4 +83,19 @@ PYBIND11_MODULE(sim, m, pybind11::mod_gil_not_used()) {
.def("set_noise_parameters", &StepControl::SetNoiseParameters,
py::arg("noise_scale"), py::arg("noise_rate"),
"Sets the noise parameters.");
using SimHistory = mujoco::platform::SimHistory;
constexpr int max_history = 2048;
constexpr int max_bytes = 128 * 1024 * 1024; // 128 MiB
py::class_<SimHistory>(m, "SimHistory")
.def(py::init<>())
.def("init", &SimHistory::Init, py::arg("state_size"),
py::arg("max_history") = max_history,
py::arg("max_bytes") = max_bytes,
"Clears and initializes the history buffer to hold `state_size` "
"mjtNum states.")
.def("get_index", &SimHistory::GetIndex,
"Returns the current history offset (0 is the most recent state).")
.def("size", &SimHistory::Size,
"Returns the number of recorded states.");
}
python/mujoco/experimental/studio/studio_app.py
+39
View File
@@ -101,6 +101,7 @@ class StudioApp:
self.model_path = model_path
self.step_control = sim.StepControl()
self.ux_state = ux.UxState()
self._setup_history()
self.status = f'Loaded: {os.path.basename(model_path)!r}'
return model, data
@@ -117,6 +118,8 @@ class StudioApp:
self.step_control = sim.StepControl()
self.ux_state = ux.UxState()
self.sim_history = sim.SimHistory()
self._setup_history()
self.theme = ux.GuiTheme.LIGHT
self.show_info = False
self.show_solver = False
@@ -209,6 +212,27 @@ class StudioApp:
"""Reset the physics."""
mujoco.mj_resetData(self.model, self.data)
mujoco.mj_forward(self.model, self.data)
self._setup_history()
def _setup_history(self) -> None:
"""(Re)initialize simulation history recording for the current model."""
ux.setup_history(
self.step_control, self.sim_history, self.ux_state, self.model,
self.data,
)
def reload_model(self) -> None:
"""Reload the current model from its file."""
if self.model_path:
self.load_model_from_file(self.model_path)
def align_camera(self, camera: mujoco.MjvCamera) -> None:
"""Recenter the camera on the model's home camera, else the free camera."""
cam_id = self.model.vis.global_.cameraid
if 0 <= cam_id < self.model.ncam:
self.ux_state.camera_index = ux.set_camera(self.model, camera, cam_id)
else:
mujoco.mjv_defaultFreeCamera(self.model, camera)
def apply_perturb(self, perturb: mujoco.MjvPerturb) -> None:
"""Apply perturbation the model."""
@@ -236,6 +260,10 @@ class StudioApp:
called instead of ``step_control.advance``. The function receives
``(model, data)`` and should step the simulation in-place.
"""
if self.ux_state.update_threadpool:
mujoco.mju_threadpool(self.data, self.ux_state.nthread)
self.ux_state.update_threadpool = False
self.apply_perturb(perturb)
if step_fn is not None:
@@ -389,6 +417,17 @@ class StudioApp:
)
imgui.Begin('Options')
# The Simulation panel draws its own collapsible section header.
ux.simulation_gui(
self.model,
self.data,
self.step_control,
self.sim_history,
self.ux_state,
self.reset_physics,
self.reload_model,
lambda: self.align_camera(camera),
)
if imgui.TreeNodeEx('Physics Settings', node_flags):
ux.physics_gui(self.model)
imgui.TreePop()
python/mujoco/experimental/studio/ux.cc
+79
View File
@@ -15,6 +15,7 @@
// Python bindings for MuJoCo platform UX components.
#include <array>
#include <span>
#include <string>
#include <tuple>
#include <vector>
@@ -23,6 +24,7 @@
#include <implot.h>
#include <mujoco/mujoco.h>
#include <mujoco/experimental/platform/helpers.h>
#include <mujoco/experimental/platform/sim/sim_history.h>
#include <mujoco/experimental/platform/sim/step_control.h>
#include <mujoco/experimental/platform/ux/gui.h>
#include <mujoco/experimental/platform/ux/interaction.h>
@@ -46,6 +48,12 @@ struct UxState {
// Read/edited by camera_selection_gui
int camera_index = mujoco::platform::kTumbleCameraIdx;
// Read/edited by simulation_gui.
int key_idx = 0;
int nthread = 0;
bool update_threadpool = false;
mujoco::platform::SimulationTimelineState timeline;
};
struct RenderFlags {
@@ -72,6 +80,9 @@ PYBIND11_MODULE(ux, m, pybind11::mod_gil_not_used()) {
.def_readwrite("state_sig", &UxState::state_sig)
.def_readwrite("watch_field_index", &UxState::watch_field_index)
.def_readwrite("camera_index", &UxState::camera_index)
.def_readwrite("key_idx", &UxState::key_idx)
.def_readwrite("nthread", &UxState::nthread)
.def_readwrite("update_threadpool", &UxState::update_threadpool)
.def_property(
"watch_field_name",
[](const UxState& self) {
@@ -128,6 +139,74 @@ PYBIND11_MODULE(ux, m, pybind11::mod_gil_not_used()) {
"Render the simulation stepping control GUI. Modifies "
"ux_state.speed_index.");
m.def(
"setup_history",
[](mujoco::platform::StepControl* step_control,
mujoco::platform::SimHistory* history, UxState& ux_state,
py::object model_obj, py::object data_obj) {
mjModel* model =
py::cast<mujoco::python::MjModelWrapper&>(model_obj).get();
mjData* data = py::cast<mujoco::python::MjDataWrapper&>(data_obj).get();
mujoco::platform::SimulationTimelineState* timeline =
&ux_state.timeline;
py::gil_scoped_release no_gil;
// Record every simulation step into the history buffer (in C++, so no
// Python is called per step). Matches the native Studio app.
history->Init(mj_stateSize(model, mjSTATE_INTEGRATION));
step_control->SetPostStepCallback(
[history, timeline](const mjModel* m, mjData* d) {
std::span<mjtNum> state = history->AddToHistory();
if (!state.empty()) {
mj_getState(m, d, state.data(), mjSTATE_INTEGRATION);
timeline->sim_head_time = d->time;
}
});
// Record the initial state and reset the scrubber.
std::span<mjtNum> state = history->AddToHistory();
if (!state.empty()) {
mj_getState(model, data, state.data(), mjSTATE_INTEGRATION);
}
*timeline = {};
timeline->sim_head_time = data->time;
},
py::arg("step_control"), py::arg("history"), py::arg("ux_state"),
py::arg("model"), py::arg("data"),
"Wire history recording: (re)initialize the buffer, install a per-step "
"recorder on step_control, record the current state and reset the "
"timeline. Call on model load and after a reset.");
m.def(
"simulation_gui",
[](py::object model_obj, py::object data_obj,
mujoco::platform::StepControl* step_control,
mujoco::platform::SimHistory* history, UxState& ux_state,
py::function reset, py::function reload, py::function align) {
mjModel* model =
py::cast<mujoco::python::MjModelWrapper&>(model_obj).get();
mjData* data = py::cast<mujoco::python::MjDataWrapper&>(data_obj).get();
// The GIL is held throughout: the callbacks call back into Python.
mujoco::platform::SimulationGuiContext ctx;
ctx.model = model;
ctx.data = data;
ctx.step_control = step_control;
ctx.history = history;
ctx.timeline = &ux_state.timeline;
ctx.speed_index = &ux_state.speed_index;
ctx.key_idx = &ux_state.key_idx;
ctx.nthread = &ux_state.nthread;
ctx.update_threadpool = &ux_state.update_threadpool;
ctx.reset = [&reset]() { reset(); };
ctx.reload = [&reload]() { reload(); };
ctx.align = [&align]() { align(); };
mujoco::platform::SimulationGui(ctx);
},
py::arg("model"), py::arg("data"), py::arg("step_control"),
py::arg("history"), py::arg("ux_state"), py::arg("reset"),
py::arg("reload"), py::arg("align"),
"Render the full Simulation panel: reset/reload/align, run/pause, speed, "
"the history scrubber, keyframes and thread count. The three callbacks "
"are invoked for the corresponding buttons.");
m.def(
"theme_select_gui",
[](mujoco::platform::GuiTheme theme) {