Allow custom step function in StepControl::Advance.

The StepControl::Advance method now accepts an optional std::function to be used instead of mj_step. This allows for custom simulation logic to be executed within the stepping loop. The Python bindings for StepControl::advance have been updated to support passing a Python callable as the custom step function.

PiperOrigin-RevId: 925393132
Change-Id: Id2820c4b55cbdfee6b01adc2395e99a990d4e7a3
This commit is contained in:
Matija Kecman
2026-06-02 09:19:30 -07:00
committed by Copybara-Service
parent 3a5626bc27
commit 0e4749501c
3 changed files with 25 additions and 7 deletions
+14 -4
View File
@@ -42,11 +42,21 @@ PYBIND11_MODULE(sim, m) {
.def(py::init<>())
.def(
"advance",
[](StepControl& self, mujoco::python::MjModelWrapper& model,
mujoco::python::MjDataWrapper& data) {
return self.Advance(model.get(), data.get());
[](StepControl& self, py::object model_obj, py::object data_obj,
py::object step_fn) {
auto& model = py::cast<mujoco::python::MjModelWrapper&>(model_obj);
auto& data = py::cast<mujoco::python::MjDataWrapper&>(data_obj);
if (step_fn.is_none()) {
return self.Advance(model.get(), data.get());
} else {
return self.Advance(
model.get(), data.get(),
[step_fn, model_obj, data_obj](mjModel*, mjData*) {
step_fn(model_obj, data_obj);
});
}
},
py::arg("model"), py::arg("data"),
py::arg("model"), py::arg("data"), py::arg("step_fn") = py::none(),
"Step physics forward, respecting speed settings and refresh budget.")
.def("force_sync", &StepControl::ForceSync,
"Ensures the next Advance() will synchronize time and step once.")