WIP Python bindings for Simulate tested on Linux only

This commit is contained in:
Levi Burner
2022-03-30 12:51:56 -04:00
parent 5ef8cb6bbc
commit 0b6bf42840
8 changed files with 517 additions and 6 deletions
+51 -2
View File
@@ -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
"$<TARGET_FILE:_callbacks>"
"$<TARGET_FILE:_constants>"
@@ -378,8 +423,10 @@ set(LIBRARIES_FOR_WHEEL
"$<TARGET_FILE:_functions>"
"$<TARGET_FILE:_render>"
"$<TARGET_FILE:_rollout>"
"$<TARGET_FILE:_simulate>"
"$<TARGET_FILE:_structs>"
"$<TARGET_FILE:mujoco>"
"$<TARGET_FILE:mjsimulate>"
)
if(MUJOCO_PYTHON_MAKE_WHEEL)
@@ -405,7 +452,9 @@ if(MUJOCO_PYTHON_MAKE_WHEEL)
_functions
_render
_rollout
_simulate
_structs
mujoco
mjsimulate
)
endif()
+4
View File
@@ -15,6 +15,7 @@
#include <utility>
#include <vector>
#include <mujoco/mjmodel.h>
#include <mujoco/simulate.h>
#include <mujoco/mjvisualize.h>
#include <mujoco/mujoco.h>
#include <pybind11/cast.h>
@@ -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);
+4
View File
@@ -19,6 +19,7 @@
#include <mujoco/mjmodel.h>
#include <mujoco/mjrender.h>
#include <mujoco/mjvisualize.h>
#include <mujoco/simulate.h>
// 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_
+142
View File
@@ -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 <array>
#include <cstdio>
#include <iostream>
#include <optional>
#include <sstream>
#include <string>
#include <mujoco/mujoco.h>
#include <mujoco/simulate.h>
#include "raw.h"
#include "structs.h"
#include <pybind11/buffer_info.h>
#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
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_<mujoco::raw::Simulate>(pymodule, "Simulate")
.def(py::init<>())
.def("renderloop",
[](mujoco::raw::Simulate& simulate) {
simulate.renderloop();
},
py::call_guard<py::gil_scoped_release>())
.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<py::gil_scoped_release>())
.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<py::gil_scoped_release>())
.def("unlock",
[](mujoco::raw::Simulate& simulate) {
simulate.mtx.unlock();
},
py::call_guard<py::gil_scoped_release>())
.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
}
+256
View File
@@ -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()
+5 -4
View File
@@ -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*',
]),
},
)
+25
View File
@@ -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.
+30
View File
@@ -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()