Expose a handle for the Python viewer.
This change also requires user scripts to explicitly synchronize changes to physics state to the viewer. The Simulate class was reconfigured so that certain UI events are handled during this sync operation, outside of the render loop on the main thread. These correspond to operations that require access to the full mjModel/mjData. To support other, more interactive operations (e.g. camera movements), a new mjvSceneState struct is introduced which captures only the portion of the physics state required for scene re-rendering. The mjvSceneState is updated from mjModel/mjData during the viewer sync operation, and is significantly cheaper than a full mj_copyModel and mj_copyData. Fixes https://github.com/deepmind/mujoco/issues/796 PiperOrigin-RevId: 525723636 Change-Id: Id08d0210a2c067d5afe85e2bf104f276aeddd75e
This commit is contained in:
Saran Tunyasuvunakool
committed by
Copybara-Service
Copybara-Service
parent
4f5da9c554
commit
b362cb4972
+77
-39
@@ -12,29 +12,76 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include <cstdint>
|
||||
#include <atomic>
|
||||
#include <cstring>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
|
||||
#include <glfw_adapter.h>
|
||||
#include <glfw_dispatch.h>
|
||||
#include <simulate.h>
|
||||
#include "structs.h"
|
||||
#include <pybind11/gil.h>
|
||||
#include <pybind11/pybind11.h>
|
||||
#include <pybind11/pytypes.h>
|
||||
|
||||
namespace mujoco::python {
|
||||
namespace {
|
||||
namespace py = ::pybind11;
|
||||
|
||||
template <typename T, int N>
|
||||
constexpr inline std::size_t sizeof_arr(const T(&arr)[N]) {
|
||||
return sizeof(arr);
|
||||
}
|
||||
|
||||
PYBIND11_MODULE(_simulate, pymodule) {
|
||||
namespace py = ::pybind11;
|
||||
using SimulateMutex = decltype(mujoco::Simulate::mtx);
|
||||
class SimulateWrapper : public mujoco::Simulate {
|
||||
public:
|
||||
SimulateWrapper(std::unique_ptr<PlatformUIAdapter> platform_ui_adapter,
|
||||
py::object scn, py::object cam,
|
||||
py::object opt, py::object pert, bool fully_managed)
|
||||
: Simulate(std::move(platform_ui_adapter),
|
||||
scn.cast<MjvSceneWrapper&>().get(),
|
||||
cam.cast<MjvCameraWrapper&>().get(),
|
||||
opt.cast<MjvOptionWrapper&>().get(),
|
||||
pert.cast<MjvPerturbWrapper&>().get(),
|
||||
fully_managed),
|
||||
m_(py::none()),
|
||||
d_(py::none()),
|
||||
scn_(scn),
|
||||
cam_(cam),
|
||||
opt_(opt),
|
||||
pert_(pert) {}
|
||||
|
||||
py::class_<SimulateMutex>(pymodule, "SimulateMutex")
|
||||
void Load(py::object m, py::object d, const std::string& path) {
|
||||
mjModel* m_raw = m.cast<MjModelWrapper&>().get();
|
||||
mjData* d_raw = d.cast<MjDataWrapper&>().get();
|
||||
{
|
||||
py::gil_scoped_release no_gil;
|
||||
Simulate::Load(m_raw, d_raw, path.c_str());
|
||||
}
|
||||
m_ = m;
|
||||
d_ = d;
|
||||
m_raw_ = m_raw;
|
||||
d_raw_ = d_raw;
|
||||
}
|
||||
|
||||
private:
|
||||
// Hold references to keep these Python objects alive for as long as the
|
||||
// simulate object.
|
||||
py::object m_;
|
||||
py::object d_;
|
||||
py::object scn_;
|
||||
py::object cam_;
|
||||
py::object opt_;
|
||||
py::object pert_;
|
||||
|
||||
mjModel* m_raw_ = nullptr;
|
||||
mjData* d_raw_ = nullptr;
|
||||
};
|
||||
|
||||
PYBIND11_MODULE(_simulate, pymodule) {
|
||||
py::class_<SimulateMutex>(pymodule, "Mutex")
|
||||
.def(
|
||||
"__enter__", [](SimulateMutex& mtx) { mtx.lock(); },
|
||||
py::call_guard<py::gil_scoped_release>())
|
||||
@@ -45,36 +92,29 @@ PYBIND11_MODULE(_simulate, pymodule) {
|
||||
},
|
||||
py::call_guard<py::gil_scoped_release>());
|
||||
|
||||
py::class_<mujoco::Simulate>(pymodule, "Simulate")
|
||||
.def(py::init([]() {
|
||||
return std::make_unique<mujoco::Simulate>(
|
||||
std::make_unique<mujoco::GlfwAdapter>());
|
||||
py::class_<SimulateWrapper>(pymodule, "Simulate")
|
||||
.def_readonly_static("MAX_GEOM", &mujoco::Simulate::kMaxGeom)
|
||||
.def(py::init([](py::object scn, py::object cam, py::object opt,
|
||||
py::object pert, bool fully_managed) {
|
||||
return std::make_unique<SimulateWrapper>(
|
||||
std::make_unique<mujoco::GlfwAdapter>(), scn, cam, opt, pert,
|
||||
fully_managed);
|
||||
}))
|
||||
.def(
|
||||
"render_loop",
|
||||
[](mujoco::Simulate& simulate) { simulate.RenderLoop(); },
|
||||
py::call_guard<py::gil_scoped_release>())
|
||||
.def(
|
||||
"load",
|
||||
[](mujoco::Simulate& simulate, MjModelWrapper& m, MjDataWrapper& d,
|
||||
const std::string& path) {
|
||||
simulate.Load(m.get(), d.get(), path.c_str());
|
||||
},
|
||||
py::call_guard<py::gil_scoped_release>())
|
||||
.def("apply_pose_perturbations",
|
||||
&mujoco::Simulate::ApplyPosePerturbations,
|
||||
py::call_guard<py::gil_scoped_release>())
|
||||
.def("apply_force_perturbations",
|
||||
&mujoco::Simulate::ApplyForcePerturbations,
|
||||
.def("load", &SimulateWrapper::Load)
|
||||
.def("sync", &mujoco::Simulate::Sync,
|
||||
py::call_guard<py::gil_scoped_release>())
|
||||
|
||||
.def(
|
||||
"render_loop",
|
||||
[](SimulateWrapper& simulate) { simulate.RenderLoop(); },
|
||||
py::call_guard<py::gil_scoped_release>())
|
||||
.def(
|
||||
"lock",
|
||||
[](mujoco::Simulate& simulate) -> SimulateMutex& {
|
||||
[](SimulateWrapper& simulate) -> SimulateMutex& {
|
||||
return simulate.mtx;
|
||||
},
|
||||
py::call_guard<py::gil_scoped_release>(),
|
||||
py::return_value_policy::reference)
|
||||
py::return_value_policy::reference_internal)
|
||||
.def_readonly("ctrl_noise_std", &mujoco::Simulate::ctrl_noise_std,
|
||||
py::call_guard<py::gil_scoped_release>())
|
||||
.def_readonly("ctrl_noise_rate", &mujoco::Simulate::ctrl_noise_rate,
|
||||
@@ -96,54 +136,52 @@ PYBIND11_MODULE(_simulate, pymodule) {
|
||||
|
||||
.def_property(
|
||||
"exitrequest",
|
||||
[](mujoco::Simulate& simulate) {
|
||||
return simulate.exitrequest.load();
|
||||
},
|
||||
[](mujoco::Simulate& simulate, bool exitrequest) {
|
||||
[](SimulateWrapper& simulate) { return simulate.exitrequest.load(); },
|
||||
[](SimulateWrapper& simulate, int exitrequest) {
|
||||
simulate.exitrequest.store(exitrequest);
|
||||
},
|
||||
py::call_guard<py::gil_scoped_release>())
|
||||
|
||||
.def_property_readonly(
|
||||
"uiloadrequest",
|
||||
[](mujoco::Simulate& simulate) {
|
||||
[](SimulateWrapper& simulate) {
|
||||
return simulate.uiloadrequest.load();
|
||||
},
|
||||
py::call_guard<py::gil_scoped_release>())
|
||||
.def(
|
||||
"uiloadrequest_decrement",
|
||||
[](mujoco::Simulate& simulate) {
|
||||
[](SimulateWrapper& simulate) {
|
||||
simulate.uiloadrequest.fetch_sub(1);
|
||||
},
|
||||
py::call_guard<py::gil_scoped_release>())
|
||||
|
||||
.def_property(
|
||||
"droploadrequest",
|
||||
[](mujoco::Simulate& simulate) {
|
||||
[](SimulateWrapper& simulate) {
|
||||
return simulate.droploadrequest.load();
|
||||
},
|
||||
[](mujoco::Simulate& simulate, bool droploadrequest) {
|
||||
[](SimulateWrapper& simulate, bool droploadrequest) {
|
||||
simulate.droploadrequest.store(droploadrequest);
|
||||
},
|
||||
py::call_guard<py::gil_scoped_release>())
|
||||
.def_property_readonly(
|
||||
"dropfilename",
|
||||
[](mujoco::Simulate& simulate) -> std::string {
|
||||
[](SimulateWrapper& simulate) -> std::string {
|
||||
return simulate.dropfilename;
|
||||
},
|
||||
py::call_guard<py::gil_scoped_release>())
|
||||
.def_property_readonly(
|
||||
"filename",
|
||||
[](mujoco::Simulate& simulate) -> std::string {
|
||||
[](SimulateWrapper& simulate) -> std::string {
|
||||
return simulate.filename;
|
||||
},
|
||||
py::call_guard<py::gil_scoped_release>())
|
||||
.def_property(
|
||||
"load_error",
|
||||
[](mujoco::Simulate& simulate) -> std::string {
|
||||
[](SimulateWrapper& simulate) -> std::string {
|
||||
return simulate.load_error;
|
||||
},
|
||||
[](mujoco::Simulate& simulate, const std::string& error) {
|
||||
[](SimulateWrapper& simulate, const std::string& error) {
|
||||
const auto max_length = sizeof_arr(simulate.load_error);
|
||||
std::strncpy(simulate.load_error, error.c_str(), max_length - 1);
|
||||
simulate.load_error[max_length - 1] = '\0';
|
||||
|
||||
Reference in New Issue
Block a user