From 0b6bf42840552e0100a375adceb2581ff6e5a30a Mon Sep 17 00:00:00 2001 From: Levi Burner Date: Wed, 30 Mar 2022 12:51:56 -0400 Subject: [PATCH] WIP Python bindings for Simulate tested on Linux only --- python/mujoco/CMakeLists.txt | 53 +++++++- python/mujoco/constants.cc | 4 + python/mujoco/raw.h | 4 + python/mujoco/simulate.cc | 142 +++++++++++++++++++ python/mujoco/simulate.py | 256 +++++++++++++++++++++++++++++++++++ python/setup.py | 9 +- simulate/CMakeLists.txt | 25 ++++ simulate/main.py | 30 ++++ 8 files changed, 517 insertions(+), 6 deletions(-) create mode 100644 python/mujoco/simulate.cc create mode 100644 python/mujoco/simulate.py create mode 100644 simulate/main.py diff --git a/python/mujoco/CMakeLists.txt b/python/mujoco/CMakeLists.txt index 4bf4b5cf..dbe7b1fc 100644 --- a/python/mujoco/CMakeLists.txt +++ b/python/mujoco/CMakeLists.txt @@ -62,6 +62,7 @@ add_compile_options("${MUJOCO_HARDEN_COMPILE_OPTIONS}") add_link_options("${MUJOCO_HARDEN_LINK_OPTIONS}") find_package(Python3 COMPONENTS Interpreter Development) +find_package(glfw3 3.3 REQUIRED) include(FindOrFetch) @@ -115,6 +116,40 @@ if(NOT TARGET mujoco) endif() endif() +# ==================== MJSIMULATE LIBRARY ========================================== +if(NOT TARGET mjsimulate) + find_library(MJSIMULATE_LIBRARY mjsimulate mjsimulate HINTS ${MUJOCO_LIBRARY_DIR} REQUIRED) + find_path(MJSIMULATE_INCLUDE mujoco/simulate.h mujoco/uitools.h mujoco/array_safety.h HINTS ${MUJOCO_INCLUDE_DIR} REQUIRED) + message("MuJoCo Simulate is at ${MJSIMULATE_LIBRARY}") + message("MuJoCo Simulate headers are at ${MJSIMULATE_INCLUDE}") + add_library(mjsimulate SHARED IMPORTED) + if(WIN32) + set_target_properties(mjsimulate PROPERTIES IMPORTED_IMPLIB "${MJSIMULATE_LIBRARY}") + else() + set_target_properties(mjsimulate PROPERTIES IMPORTED_LOCATION "${MJSIMULATE_LIBRARY}") + endif() + target_include_directories(mjsimulate INTERFACE "${MJSIMULATE_INCLUDE}") + if(APPLE) + execute_process( + COMMAND otool -XD ${MJSIMULATE_LIBRARY} + COMMAND head -n 1 + COMMAND xargs dirname + COMMAND xargs echo -n + OUTPUT_VARIABLE MJSIMULATE_INSTALL_NAME_DIR + ) + set_target_properties(mjsimulate PROPERTIES INSTALL_NAME_DIR "${MJSIMULATE_INSTALL_NAME_DIR}") + elseif(UNIX) + execute_process( + COMMAND objdump -p ${MJSIMULATE_LIBRARY} + COMMAND grep SONAME + COMMAND grep -Po [^\\s]+$ + COMMAND xargs echo -n + OUTPUT_VARIABLE MJSIMULATE_SONAME + ) + set_target_properties(mjsimulate PROPERTIES IMPORTED_SONAME "${MJSIMULATE_SONAME}") + endif() +endif() + # ==================== ABSEIL ================================================== set(MUJOCO_PYTHON_ABSL_TARGETS absl::core_headers absl::flat_hash_map absl::span) findorfetch( @@ -231,7 +266,7 @@ target_link_libraries(errors_header INTERFACE crossplatform func_wrap mujoco) add_library(raw INTERFACE) target_sources(raw INTERFACE raw.h) set_target_properties(raw PROPERTIES PUBLIC_HEADER raw.h) -target_link_libraries(raw INTERFACE mujoco) +target_link_libraries(raw INTERFACE mujoco mjsimulate) add_library(structs_header INTERFACE) target_sources( @@ -309,7 +344,7 @@ target_link_libraries( ) mujoco_pybind11_module(_constants constants.cc) -target_link_libraries(_constants PRIVATE mujoco) +target_link_libraries(_constants PRIVATE mujoco mjsimulate) mujoco_pybind11_module(_enums enums.cc) target_link_libraries( @@ -370,6 +405,16 @@ target_link_libraries( structs_header ) +mujoco_pybind11_module(_simulate simulate.cc) +target_link_libraries( + _simulate + PRIVATE mjsimulate + mujoco + raw + glfw + structs_header) +target_link_options(_simulate PRIVATE -Wl,-no-as-needed) + set(LIBRARIES_FOR_WHEEL "$" "$" @@ -378,8 +423,10 @@ set(LIBRARIES_FOR_WHEEL "$" "$" "$" + "$" "$" "$" + "$" ) if(MUJOCO_PYTHON_MAKE_WHEEL) @@ -405,7 +452,9 @@ if(MUJOCO_PYTHON_MAKE_WHEEL) _functions _render _rollout + _simulate _structs mujoco + mjsimulate ) endif() diff --git a/python/mujoco/constants.cc b/python/mujoco/constants.cc index 2ffe73bb..3fcbd279 100644 --- a/python/mujoco/constants.cc +++ b/python/mujoco/constants.cc @@ -15,6 +15,7 @@ #include #include #include +#include #include #include #include @@ -78,6 +79,9 @@ PYBIND11_MODULE(_constants, pymodule) { // from mujoco.h X(mjVERSION_HEADER); + // from simulate.h + X(mujoco::Simulate::kMaxFilenameLength); + #undef X pymodule.attr("mjDISABLESTRING") = MakeTuple(mjDISABLESTRING); pymodule.attr("mjENABLESTRING") = MakeTuple(mjENABLESTRING); diff --git a/python/mujoco/raw.h b/python/mujoco/raw.h index 5aa146f4..3399540a 100644 --- a/python/mujoco/raw.h +++ b/python/mujoco/raw.h @@ -19,6 +19,7 @@ #include #include #include +#include // Type aliases for MuJoCo C structs to allow us refer to consistently refer // to them under the "raw" namespace. @@ -55,6 +56,9 @@ using MjvOption = ::mjvOption; using MjvScene = ::mjvScene; using MjvFigure = ::mjvFigure; +// From simulate.h +using Simulate = ::mujoco::Simulate; + } // namespace mujoco::raw #endif // MUJOCO_PYTHON_RAW_H_ diff --git a/python/mujoco/simulate.cc b/python/mujoco/simulate.cc new file mode 100644 index 00000000..46540e8b --- /dev/null +++ b/python/mujoco/simulate.cc @@ -0,0 +1,142 @@ +// Copyright 2022 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include + +#include +#include +#include "raw.h" +#include "structs.h" +#include +#include +#include +#include + +namespace mujoco::python { + +namespace { + +namespace py = ::pybind11; + +const auto simulate_doc = R"( +Python wrapper for the Simulate class +)"; + +// We define SimulateWrapper here instead of in structs because +// we do not want to make _structs dependent on gflw + +PYBIND11_MODULE(_simulate, pymodule) { + namespace py = ::pybind11; + + py::class_(pymodule, "Simulate") + .def(py::init<>()) + .def("renderloop", + [](mujoco::raw::Simulate& simulate) { + simulate.renderloop(); + }, + py::call_guard()) + .def("load", + [](mujoco::raw::Simulate& simulate, std::string filename, const MjModelWrapper& m, MjDataWrapper& d, bool delete_old_m_d) { + const raw::MjModel* m_ptr = m.get(); + raw::MjData* d_ptr = d.get(); + simulate.load(filename.c_str(), (mjModel*)m_ptr, d_ptr, delete_old_m_d); + }, + py::call_guard()) + .def("applyposepertubations", &mujoco::raw::Simulate::applyposepertubations) + .def("applyforceperturbations", &mujoco::raw::Simulate::applyforceperturbations) + + .def("lock", // TODO wrap mutex properly as as seperate pybind11 object? + [](mujoco::raw::Simulate& simulate) { + simulate.mtx.lock(); + }, + py::call_guard()) + .def("unlock", + [](mujoco::raw::Simulate& simulate) { + simulate.mtx.unlock(); + }, + py::call_guard()) + .def_readwrite("ctrlnoisestd", &mujoco::raw::Simulate::ctrlnoisestd) + .def_readwrite("ctrlnoiserate", &mujoco::raw::Simulate::ctrlnoiserate) + .def_readwrite("slow_down", &mujoco::raw::Simulate::slow_down) + .def_readwrite("speed_changed", &mujoco::raw::Simulate::speed_changed) + .def("getrefreshRate", + [](mujoco::raw::Simulate& simulate) { + return simulate.vmode.refreshRate; + }) + + .def_readwrite("busywait", &mujoco::raw::Simulate::busywait) + .def_readwrite("run", &mujoco::raw::Simulate::run) + //.def_readwrite("exitrequest", &mujoco::raw::Simulate::exitrequest) + .def("getexitrequest", + [](mujoco::raw::Simulate& simulate) { + return simulate.exitrequest.load(); + } + ) + .def("setexitrequest", + [](mujoco::raw::Simulate& simulate, bool exitrequest) { + simulate.exitrequest.store(exitrequest); + } + ) + // .def_readwrite("uiloadrequest", &mujoco::raw::Simulate::uiloadrequest) + .def("getuiloadrequest", + [](mujoco::raw::Simulate& simulate) { + return simulate.uiloadrequest.load(); + } + ) + .def("setuiloadrequest", + [](mujoco::raw::Simulate& simulate, int uiloadrequest) { + simulate.uiloadrequest.store(uiloadrequest); + } + ) + .def("uiloadrequest_fetch_sub", + [](mujoco::raw::Simulate& simulate, int arg) { + simulate.uiloadrequest.fetch_sub(arg); + } + ) + // .def_readwrite("droploadrequest", &mujoco::raw::Simulate::droploadrequest) + .def("getdroploadrequest", + [](mujoco::raw::Simulate& simulate) { + return simulate.droploadrequest.load(); + } + ) + .def("setdroploadrequest", + [](mujoco::raw::Simulate& simulate, bool droploadrequest) { + simulate.droploadrequest.store(droploadrequest); + } + ) + .def("getdropfilename", + [](mujoco::raw::Simulate& simulate) { + return (char*)simulate.dropfilename; + } + ) + .def("getfilename", + [](mujoco::raw::Simulate& simulate) { + return (char*)simulate.filename; + } + ) + .def("setloadError", + [](mujoco::raw::Simulate& simulate, std::string& loadError) { + strncpy(simulate.loadError, loadError.c_str(), simulate.kMaxFilenameLength); + } + ); +} + +} // namespace + +} diff --git a/python/mujoco/simulate.py b/python/mujoco/simulate.py new file mode 100644 index 00000000..d4c37f14 --- /dev/null +++ b/python/mujoco/simulate.py @@ -0,0 +1,256 @@ +# Copyright 2022 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Wrap the pybind11 bound Simulate class to provide type conversions""" + +import mujoco +from mujoco import _simulate + +import glfw +import numpy as np + +import ctypes +import math +import threading +import time + +class Simulate(_simulate.Simulate): + def __init__(self): + super().__init__() + + def _get_refreshRate(self): + return self.getrefreshRate() + + refreshRate = property( + fget=_get_refreshRate) + + def _get_drop_file_name(self): + return self.getdropfilename() + + dropfilename = property( + fget=_get_drop_file_name, + doc="char* from Simulate" + ) + + def _get_file_name(self): + return self.getfilename() + + filename = property( + fget=_get_file_name, + doc="char* from Simulate" + ) + + def _set_loadError(self, loadError): + self.setloadError(loadError) + + loadError = property( + fset=_set_loadError, + doc="set string to Simulate" + ) + + def _get_exit_request(self): + return self.getexitrequest() + + def _set_exit_request(self, exitrequest): + self.setexitrequest(exitrequest) + + exitrequest = property( + fget=_get_exit_request, + fset=_set_exit_request, + doc="atomic bool from Simulate") + + def _get_ui_load_request(self): + return self.getuiloadrequest() + + def _set_ui_load_request(self, uiloadrequest): + self.setuiloadrequest(uiloadrequest) + + uiloadrequest = property( + fget=_get_ui_load_request, + fset=_set_ui_load_request, + doc="atomic int from Simulate") + + def _get_drop_load_request(self): + return self.getdroploadrequest() + + def _set_drop_load_request(self, droploadrequest): + self.setdroploadrequest(droploadrequest) + + droploadrequest = property( + fget=_get_drop_load_request, + fset=_set_drop_load_request, + doc="atomic bool from Simulate") + +def load_and_step_model(m, d, simulate, filename, preload_callback=None, load_callback=None): + if preload_callback is not None: + preload_callback(m, d) # Call with old model/data so user can do cleanup + + try: + mnew = mujoco.MjModel.from_xml_path(filename) + loadError = None + except Exception as e: + print('Error loading using from_xml_path') + print(e) + mnew = None + loadError = str(e) + + if mnew is not None: + dnew = mujoco.MjData(mnew) + simulate.load(filename, mnew, dnew, False) + mujoco.mj_forward(mnew, dnew) + else: + mnew = None + dnew = None + + if load_callback is not None: + load_callback(mnew, dnew, loadError) + + return mnew, dnew, loadError + +def run_physics_loop(simulate, preload_callback=None, load_callback=None, file=None): + # request loadmodel if file given (otherwise drag-and-drop) + if file is not None: + m, d, loadError = load_and_step_model(None, None, simulate, file, preload_callback, load_callback) + ctrlnoise = np.zeros((m.nu,)) + if loadError is not None: + simulate.loadError = loadError + else: + m = None + d = None + ctrlnoise = None + + # constants + syncmisalign = 0.1 # maximum time mis-alignment before re-sync + refreshfactor = 0.7 # fraction of refresh available for simulation + + # cpu-sim synchronization point + cpusync = 0.0 + simsync = 0.0 + + # run until asked to exit + while not simulate.exitrequest: + if simulate.droploadrequest: + simulate.droploadrequest = 0 + m_, d_, loadError = load_and_step_model(m, d, simulate, simulate.dropfilename, + preload_callback, load_callback) + + if m_ is not None: + m = m_ + d = d_ + ctrlnoise = np.zeros((m.nu,)) + else: + simulate.loadError = loadError + + if simulate.uiloadrequest: + simulate.uiloadrequest_fetch_sub(1) + m_, d_, loadError = load_and_step_model(m, d, simulate, simulate.filename, + preload_callback, load_callback) + if m_ is not None: + m = m_ + d = d_ + ctrlnoise = np.zeros((m.nu,)) + else: + simulate.loadError = loadError + + # sleep for 1 ms or yield, to let main thread run + # yield results in busy wait - which has better timing but kills battery life + if simulate.run != 0 and simulate.busywait != 0: + time.sleep(0) + else: + time.sleep(0.001) + + # Start exclusive access + simulate.lock() + + # run only if model is present + if m is not None: + # running + if simulate.run != 0: + # record cpu time at start of iteration + tmstart = glfw.get_time() + + # inject noise + if simulate.ctrlnoisestd != 0.0: + # convert rate and scale to discrete time given current timestep + rate = math.exp(-m.opt.timestep / simulate.ctrlnoiserate) + scale = simulate.ctrlnoisestd * math.sqrt(1-rate*rate) + + for i in range(m.nu): + # update noise + ctrlnoise[i] = rate * ctrlnoise[i] + scale * mujoco.mju_standardNormal(None) + # apply noise + d.ctrl[i] = ctrlnoise[i] + + # out-of-sync (for any reason) + offset = abs((d.time*simulate.slow_down - simsync) - (tmstart - cpusync)) + if (d.time*simulate.slow_down < simsync or tmstart < cpusync or cpusync == 0.0 or + offset > syncmisalign*simulate.slow_down or simulate.speed_changed): + # re-sync + cpusync = tmstart + simsync = d.time*simulate.slow_down + simulate.speed_changed = False + + # clear old perturbations, apply new + d.xfrc_applied[:, :] = 0 + simulate.applyposepertubations(0) # move mocap bodies only + simulate.applyforceperturbations() + + # run single step, let next iteration deal with timing + mujoco.mj_step(m, d) + + # in-sync + else: + while ((d.time*simulate.slow_down - simsync) < (glfw.get_time() - cpusync) and + (glfw.get_time() - tmstart) < (refreshfactor/simulate.refreshRate)): + # clear old perturbations, apply new + d.xfrc_applied[:, :] = 0 + simulate.applyposepertubations(0) # move mocap bodies only + simulate.applyforceperturbations() + + # run mj_step + prevtm = d.time*simulate.slow_down + mujoco.mj_step(m, d) + + # break on reset + if d.time*simulate.slow_down < prevtm: + break + + # paused + else: + # apply pose perturbation + simulate.applyposepertubations(1) # move mocap and dynamic bodies + + # run mj_forward, to update rendering and joint sliders + mujoco.mj_forward(m, d) + + # end exclusive access + simulate.unlock() + +def run_simulate_and_physics(file=None, preload_callback=None, load_callback=None): + # simulate object encapsulates the UI + simulate = Simulate() + + # init GLFW + if not glfw.init(): + raise mujoco.FatalError('could not initialize GLFW') + + # if m is not None: + physics_thread = threading.Thread(target=lambda: run_physics_loop(simulate, preload_callback=preload_callback, load_callback=load_callback, file=file)) + physics_thread.start() + + # start simulation thread (this creates the UI) + simulate.renderloop() + physics_thread.join() + + glfw.terminate() diff --git a/python/setup.py b/python/setup.py index fad9204e..e0380108 100644 --- a/python/setup.py +++ b/python/setup.py @@ -63,12 +63,11 @@ def get_mujoco_lib_pattern(): def get_external_lib_patterns(): if platform.system() == 'Windows': - return ['mujoco.dll'] + return ['mujoco.dll', 'mjsimulate.dll'] elif platform.system() == 'Darwin': - return ['libmujoco.*.dylib'] + return ['libmujoco.*.dylib', 'libmjsimulate.*.dylib'] else: - return ['libmujoco.so.*'] - + return ['libmujoco.so.*', 'libmjsimulate.so*'] def start_and_end(iterable): it = iter(iterable) @@ -292,6 +291,7 @@ setup( CMakeExtension('mujoco._functions'), CMakeExtension('mujoco._render'), CMakeExtension('mujoco._rollout'), + CMakeExtension('mujoco._simulate'), CMakeExtension('mujoco._structs'), ], python_requires='>=3.7', @@ -318,6 +318,7 @@ setup( 'libmujoco*.so.*', 'mujoco.dll', 'include/mujoco/*.h', + 'libmjsimulate*.so*', ]), }, ) diff --git a/simulate/CMakeLists.txt b/simulate/CMakeLists.txt index ae02fdc8..9a74fd77 100644 --- a/simulate/CMakeLists.txt +++ b/simulate/CMakeLists.txt @@ -117,6 +117,10 @@ target_compile_options(libsimulate PUBLIC ${MUJOCO_SIMULATE_COMPILE_OPTIONS}) target_link_libraries(libsimulate PUBLIC glfw lodepng mujoco::mujoco) target_link_options(libsimulate PRIVATE ${MUJOCO_SIMULATE_LINK_OPTIONS}) +set_target_properties( + mjsimulate PROPERTIES VERSION "${mujoco_VERSION}" PUBLIC_HEADER "simulate.h" +) + if(APPLE) target_sources(libsimulate PRIVATE macos_save.mm) target_link_libraries(libsimulate PUBLIC "-framework Cocoa") @@ -219,6 +223,17 @@ if(_INSTALL_SIMULATE) MUJOCO_ENABLE_RPATH ) + target_add_rpath( + TARGETS + mjsimulate + INSTALL_DIRECTORY + "${CMAKE_INSTALL_PREFIX}/${CMAKE_INSTALL_BINDIR}" + LIB_DIRS + "${CMAKE_INSTALL_PREFIX}/${CMAKE_INSTALL_LIBDIR}" + DEPENDS + MUJOCO_ENABLE_RPATH + ) + install( TARGETS simulate EXPORT ${PROJECT_NAME} @@ -229,6 +244,16 @@ if(_INSTALL_SIMULATE) PUBLIC_HEADER DESTINATION ${CMAKE_INSTALL_INCLUDEDIR} COMPONENT simulate ) + install( + TARGETS mjsimulate + EXPORT ${PROJECT_NAME} + RUNTIME DESTINATION "${CMAKE_INSTALL_BINDIR}" COMPONENT simulate + LIBRARY DESTINATION "${CMAKE_INSTALL_LIBDIR}" COMPONENT simulate + ARCHIVE DESTINATION "${CMAKE_INSTALL_LIBDIR}" COMPONENT simulate + BUNDLE DESTINATION "${CMAKE_INSTALL_BINDIR}" COMPONENT simulate + PUBLIC_HEADER DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/mujoco COMPONENT simulate + ) + if(NOT MUJOCO_SIMULATE_USE_SYSTEM_GLFW) # We downloaded GLFW. Depending if it is a static or shared LIBRARY we might # need to install it. diff --git a/simulate/main.py b/simulate/main.py new file mode 100644 index 00000000..da1be97c --- /dev/null +++ b/simulate/main.py @@ -0,0 +1,30 @@ +# Copyright 2022 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Python version of the simulate C++ sample""" + +import mujoco +from mujoco.simulate import run_simulate_and_physics +import sys + +if __name__ == '__main__': + # print version, check compatibility + print('MuJoCo version {}'.format(mujoco.mj_versionString())) + if mujoco.mjVERSION_HEADER != mujoco.mj_version(): + raise mujoco.FatalError('Headers and library have different versions') + + if len(sys.argv) > 1: + run_simulate_and_physics(sys.argv[1]) + else: + run_simulate_and_physics()