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

Modify Python simulate so that the Reload button works.

Browse Source 
Also allow a `Callable[[], [MjModel, MjData]]` to be provided as an argument to `launch` instead of a fixed MjModel/MjData instance. This allows users to customise the behaviour of the Reload button.

Context: https://github.com/deepmind/mujoco/discussions/576#discussioncomment-4161429
PiperOrigin-RevId: 492296268
Change-Id: Ic4b797d2ea8a6967f7e3e7316f2184ea94a86d8b
This commit is contained in:
Saran Tunyasuvunakool
2022-12-01 14:11:49 -08:00
committed by Copybara-Service
parent 89185b4af2
commit 9fcf3b0c2e
2 changed files with 86 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/simulate.cc
+3 -4
View File
@@ -17,9 +17,7 @@
#include <string>
#include <simulate.h>
#include "raw.h"
#include "structs.h"
#include <pybind11/detail/common.h>
#include <pybind11/pybind11.h>
namespace mujoco::python {
@@ -47,8 +45,9 @@ PYBIND11_MODULE(_simulate, pymodule) {
py::call_guard<py::gil_scoped_release>())
.def(
"load",
[](mujoco::Simulate& simulate, MjModelWrapper& m, MjDataWrapper& d) {
simulate.load("", m.get(), d.get());
[](mujoco::Simulate& simulate, const std::string& path,
MjModelWrapper& m, MjDataWrapper& d) {
simulate.load(path.c_str(), m.get(), d.get());
},
py::call_guard<py::gil_scoped_release>())
.def("applyposepertubations", &mujoco::Simulate::applyposepertubations,
python/mujoco/simulate.py
+83 -45
View File
@@ -21,7 +21,7 @@ import math
import os
import threading
import time
from typing import Callable, Optional
from typing import Callable, Optional, Union
import glfw
import mujoco
@@ -48,35 +48,32 @@ MAX_SYNC_MISALIGN = 0.1
SIM_REFRESH_FRACTION = 0.7
CallbackType = Callable[[mujoco.MjModel, mujoco.MjData], None]
LoaderType = Callable[[], tuple[mujoco.MjModel, mujoco.MjData]]
# Loader function that also returns a file path for the GUI to display.
_LoaderWithPathType = Callable[[], tuple[mujoco.MjModel, mujoco.MjData, str]]
_InternalLoaderType = Union[LoaderType, _LoaderWithPathType]
Simulate = _simulate.Simulate
def _reload_from_file(simulate: Simulate, filename: str):
"""Loads an MJCF model into the Simulate GUI."""
try:
m = mujoco.MjModel.from_xml_path(filename)
except mujoco.FatalError as e:
m = None
simulate.load_error = str(e)
def _file_loader(path: str) -> _LoaderWithPathType:
"""Loads an MJCF model from file path."""
if m is not None:
def load(path=path) -> tuple[mujoco.MjModel, mujoco.MjData, str]:
m = mujoco.MjModel.from_xml_path(path)
d = mujoco.MjData(m)
simulate.load(m, d)
mujoco.mj_forward(m, d)
else:
d = None
return m, d, path
return m, d
return load
def _physics_loop(simulate: Simulate,
m: Optional[mujoco.MjModel],
d: Optional[mujoco.MjData]):
def _physics_loop(simulate: Simulate, loader: Optional[_InternalLoaderType]):
"""Physics loop for the Simulate GUI, to be run in a separate thread."""
ctrlnoise = None
if m is not None:
ctrlnoise = np.zeros((m.nu,))
m: mujoco.MjModel = None
d: mujoco.MjData = None
ctrlnoise = np.array([])
reload = True
# CPU-sim synchronization point.
synccpu = 0.0
@@ -86,20 +83,40 @@ def _physics_loop(simulate: Simulate,
while not simulate.exitrequest:
if simulate.droploadrequest:
simulate.droploadrequest = 0
new_m, new_d = _reload_from_file(simulate, simulate.dropfilename)
if new_m is not None:
m = new_m
d = new_d
ctrlnoise = np.zeros((m.nu,))
loader = _file_loader(simulate.dropfilename)
reload = True
if simulate.uiloadrequest:
simulate.uiloadrequest_decrement()
new_m, new_d = _reload_from_file(simulate, simulate.dropfilename)
if new_m is not None:
m = new_m
d = new_d
reload = True
if reload and loader is not None:
try:
load_tuple = loader()
except Exception as e: # pylint: disable=broad-except
simulate.load_error = str(e)
else:
# Do not assign to m and d until simulate.load is done!
# This is because simulate.load needs to clean up mjvScene and
# mjrContext. This cleanup logic requires access to the old m and d.
new_m, new_d = load_tuple[:2]
# If the loader does not raise an exception then we assume that it
# successfully created mjModel and mjData. This is specified in the type
# annotation, but we perform a runtime assertion here as well to prevent
# possible segmentation faults.
assert new_m is not None and new_d is not None
path = load_tuple[2] if len(load_tuple) == 3 else ''
simulate.load(path, new_m, new_d)
# We can now allow the old m and d to be deleted.
m, d = new_m, new_d
mujoco.mj_forward(m, d)
ctrlnoise = np.zeros((m.nu,))
reload = False
# Sleep for 1 ms or yield, to let main thread run.
if simulate.run != 0 and simulate.busywait != 0:
time.sleep(0)
@@ -186,15 +203,28 @@ def _physics_loop(simulate: Simulate,
mujoco.mj_forward(m, d)
def launch(model: Optional[mujoco.MjModel] = None,
data: Optional[mujoco.MjData] = None,
*,
run_physics_thread: bool = True) -> None:
"""Launches the Simulate GUI."""
def _launch_internal(model: Optional[mujoco.MjModel] = None,
data: Optional[mujoco.MjData] = None,
*,
run_physics_thread: bool = True,
loader: Optional[_InternalLoaderType] = None) -> None:
"""Internal API, so that the public API has more readable type annotations."""
if model is None and data is not None:
raise ValueError('mjData is specified but mjModel is not')
elif model is not None and data is None:
data = mujoco.MjData(model)
elif callable(model) and data is not None:
raise ValueError(
'mjData should not be specified when an mjModel loader is used')
elif loader is not None and model is not None:
raise ValueError('model and loader are both specified')
if loader is None and model is not None:
def _loader(m=model, d=data) -> tuple[mujoco.MjModel, mujoco.MjData]:
if d is None:
d = mujoco.MjData(m)
return m, d
loader = _loader
# The simulate object encapsulates the UI.
simulate = Simulate()
@@ -207,21 +237,30 @@ def launch(model: Optional[mujoco.MjModel] = None,
if run_physics_thread:
physics_thread = threading.Thread(
target=_physics_loop, args=(simulate, model, data))
target=_physics_loop, args=(simulate, loader))
physics_thread.start()
# Load the initial model, if one is given.
if model is not None:
t = threading.Thread(target=simulate.load, args=(model, data))
t.start()
del t
simulate.renderloop()
if run_physics_thread:
physics_thread.join()
def launch(model: Optional[mujoco.MjModel] = None,
data: Optional[mujoco.MjData] = None,
*,
run_physics_thread: bool = True,
loader: Optional[LoaderType] = None) -> None:
"""Launches the Simulate GUI."""
_launch_internal(
model, data, run_physics_thread=run_physics_thread, loader=loader)
def launch_from_path(path: str) -> None:
"""Launches the Simulate GUI from file path."""
_launch_internal(loader=_file_loader(path))
def launch_repl(model: mujoco.MjModel, data: mujoco.MjData) -> None:
"""EXPERIMENTAL FEATURE: Launches the Simulate GUI in REPL mode."""
try:
@@ -253,8 +292,7 @@ if __name__ == '__main__':
def main(argv) -> None:
del argv
if _MJCF_PATH.value is not None:
model = mujoco.MjModel.from_xml_path(os.path.expanduser(_MJCF_PATH.value))
launch(model)
launch_from_path(os.path.expanduser(_MJCF_PATH.value))
else:
launch()