Merge pull request #419 from aftersomemath:simulate-python

PiperOrigin-RevId: 487496285
Change-Id: I17d15bd2a3886e5ce1631f4eb78afa94830d08cc

python/MANIFEST.in

@@ -1,3 +1,3 @@
 include LICENSE *.md
-recursive-include mujoco *.h *.cc CMakeLists.txt
+recursive-include mujoco *.h *.cc *.mm CMakeLists.txt Simulate*.cmake
 recursive-include cmake *.cmake

python/make_sdist.sh

@@ -53,6 +53,9 @@ cp "${package_dir}"/../LICENSE .
 mkdir cmake
 cp "${package_dir}"/../cmake/*.cmake cmake
 
+# Copy over Simulate source code.
+cp -r "${package_dir}"/../simulate mujoco
+
 python setup.py sdist --formats=gztar
 tar -tf dist/mujoco-*.tar.gz
 popd

python/mujoco/CMakeLists.txt

@@ -119,6 +119,7 @@ if(NOT TARGET mujoco)
     )
     set_target_properties(mujoco PROPERTIES IMPORTED_SONAME "${MUJOCO_SONAME}")
   endif()
+  add_library(mujoco::mujoco ALIAS mujoco)
 endif()
 
 # ==================== ABSEIL ==================================================
@@ -190,6 +191,9 @@ findorfetch(
 )
 
 # ==================== MUJOCO PYTHON BINDINGS ==================================
+set(SIMULATE_BUILD_EXECUTABLE OFF)
+set(SIMULATE_GLFW_DYNAMIC_SYMBOLS ON)
+add_subdirectory(simulate)
 
 add_subdirectory(util)
 
@@ -376,6 +380,14 @@ target_link_libraries(
           structs_header
 )
 
+mujoco_pybind11_module(_simulate simulate.cc)
+target_link_libraries(
+  _simulate
+  PRIVATE mujoco
+          mujoco::libsimulate
+          raw
+          structs_header)
+
 set(LIBRARIES_FOR_WHEEL
     "$<TARGET_FILE:_callbacks>"
     "$<TARGET_FILE:_constants>"
@@ -384,6 +396,7 @@ set(LIBRARIES_FOR_WHEEL
     "$<TARGET_FILE:_functions>"
     "$<TARGET_FILE:_render>"
     "$<TARGET_FILE:_rollout>"
+    "$<TARGET_FILE:_simulate>"
     "$<TARGET_FILE:_structs>"
     "$<TARGET_FILE:mujoco>"
 )
@@ -411,6 +424,7 @@ if(MUJOCO_PYTHON_MAKE_WHEEL)
     _functions
     _render
     _rollout
+    _simulate
     _structs
     mujoco
   )

python/mujoco/simulate.cc

@@ -0,0 +1,143 @@
+// 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 <cstdint>
+#include <cstring>
+#include <string>
+
+#include <simulate.h>
+#include "raw.h"
+#include "structs.h"
+#include <pybind11/detail/common.h>
+#include <pybind11/pybind11.h>
+
+namespace mujoco::python {
+namespace {
+PYBIND11_MODULE(_simulate, pymodule) {
+  namespace py = ::pybind11;
+  using SimulateMutex = decltype(mujoco::Simulate::mtx);
+
+  py::class_<SimulateMutex>(pymodule, "SimulateMutex")
+      .def(
+          "__enter__", [](SimulateMutex& mtx) { mtx.lock(); },
+          py::call_guard<py::gil_scoped_release>())
+      .def(
+          "__exit__",
+          [](SimulateMutex& mtx, py::handle, py::handle, py::handle) {
+            mtx.unlock();
+          },
+          py::call_guard<py::gil_scoped_release>());
+
+  py::class_<mujoco::Simulate>(pymodule, "Simulate")
+      .def(py::init<>())
+      .def(
+          "renderloop",
+          [](mujoco::Simulate& simulate) { simulate.renderloop(); },
+          py::call_guard<py::gil_scoped_release>())
+      .def(
+          "load",
+          [](mujoco::Simulate& simulate, MjModelWrapper& m, MjDataWrapper& d) {
+            simulate.load("", m.get(), d.get());
+          },
+          py::call_guard<py::gil_scoped_release>())
+      .def("applyposepertubations", &mujoco::Simulate::applyposepertubations,
+           py::call_guard<py::gil_scoped_release>())
+      .def("applyforceperturbations",
+           &mujoco::Simulate::applyforceperturbations,
+           py::call_guard<py::gil_scoped_release>())
+
+      .def(
+          "lock",
+          [](mujoco::Simulate& simulate) -> SimulateMutex& {
+            return simulate.mtx;
+          },
+          py::call_guard<py::gil_scoped_release>(),
+          py::return_value_policy::reference)
+      .def_readonly("ctrlnoisestd", &mujoco::Simulate::ctrlnoisestd,
+                    py::call_guard<py::gil_scoped_release>())
+      .def_readonly("ctrlnoiserate", &mujoco::Simulate::ctrlnoiserate,
+                    py::call_guard<py::gil_scoped_release>())
+
+      .def_readonly("real_time_index", &mujoco::Simulate::realTimeIndex,
+                    py::call_guard<py::gil_scoped_release>())
+      .def_readwrite("speed_changed", &mujoco::Simulate::speedChanged,
+                     py::call_guard<py::gil_scoped_release>())
+      .def_readwrite("measured_slowdown", &mujoco::Simulate::measuredSlowdown,
+                     py::call_guard<py::gil_scoped_release>())
+      .def_readonly("refresh_rate", &mujoco::Simulate::refreshRate,
+                    py::call_guard<py::gil_scoped_release>())
+
+      .def_readonly("busywait", &mujoco::Simulate::busywait,
+                    py::call_guard<py::gil_scoped_release>())
+      .def_readonly("run", &mujoco::Simulate::run,
+                    py::call_guard<py::gil_scoped_release>())
+
+      .def_property_readonly(
+          "exitrequest",
+          [](mujoco::Simulate& simulate) {
+            return simulate.exitrequest.load();
+          },
+          py::call_guard<py::gil_scoped_release>())
+
+      .def_property_readonly(
+          "uiloadrequest",
+          [](mujoco::Simulate& simulate) {
+            return simulate.uiloadrequest.load();
+          },
+          py::call_guard<py::gil_scoped_release>())
+      .def(
+          "uiloadrequest_decrement",
+          [](mujoco::Simulate& simulate) {
+            simulate.uiloadrequest.fetch_sub(1);
+          },
+          py::call_guard<py::gil_scoped_release>())
+
+      .def_property(
+          "droploadrequest",
+          [](mujoco::Simulate& simulate) {
+            return simulate.droploadrequest.load();
+          },
+          [](mujoco::Simulate& simulate, bool droploadrequest) {
+            simulate.droploadrequest.store(droploadrequest);
+          },
+          py::call_guard<py::gil_scoped_release>())
+      .def_property_readonly(
+          "dropfilename",
+          [](mujoco::Simulate& simulate) -> std::string {
+            return simulate.dropfilename;
+          },
+          py::call_guard<py::gil_scoped_release>())
+      .def_property_readonly(
+          "filename",
+          [](mujoco::Simulate& simulate) -> std::string {
+            return simulate.filename;
+          },
+          py::call_guard<py::gil_scoped_release>())
+      .def_property(
+          "load_error",
+          [](mujoco::Simulate& simulate) -> std::string {
+            return simulate.loadError;
+          },
+          [](mujoco::Simulate& simulate, const std::string& error) {
+            std::strncpy(simulate.loadError, error.c_str(),
+                         simulate.kMaxFilenameLength);
+          });
+
+  pymodule.def("setglfwdlhandle", [](std::uintptr_t dlhandle) {
+    mujoco::setglfwdlhandle(reinterpret_cast<void*>(dlhandle));
+  });
+}
+
+}  // namespace
+}  // namespace mujoco::python

python/mujoco/simulate.py

@@ -0,0 +1,266 @@
+# 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 bindings for the Simulate GUI."""
+
+import atexit
+import code
+import inspect
+import math
+import os
+import threading
+import time
+import typing
+from typing import Callable, Optional
+
+import glfw
+import mujoco
+from mujoco import _simulate
+import numpy as np
+
+if not glfw._glfw:  # pylint: disable=protected-access
+  raise RuntimeError('GLFW dynamic library handle is not available')
+else:
+  _simulate.setglfwdlhandle(glfw._glfw._handle)  # pylint: disable=protected-access
+
+# Logarithmically spaced realtime slow-down coefficients (percent).
+PERCENT_REALTIME = (
+    100, 80, 66, 50, 40, 33, 25, 20, 16, 13,
+    10, 8, 6.6, 5, 4, 3.3, 2.5, 2, 1.6, 1.3,
+    1, 0.8, 0.66, 0.5, 0.4, 0.33, 0.25, 0.2, 0.16, 0.13,
+    0.1
+)
+
+# Maximum time mis-alignment before re-sync.
+MAX_SYNC_MISALIGN = 0.1
+
+# Fraction of refresh available for simulation.
+SIM_REFRESH_FRACTION = 0.7
+
+CallbackType = Callable[[mujoco.MjModel, mujoco.MjData], None]
+
+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)
+
+  if m is not None:
+    d = mujoco.MjData(m)
+    simulate.load(m, d)
+    mujoco.mj_forward(m, d)
+  else:
+    d = None
+
+  return m, d
+
+
+def _physics_loop(simulate: Simulate,
+                  m: Optional[mujoco.MjModel],
+                  d: Optional[mujoco.MjData]):
+  """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,))
+
+  # CPU-sim synchronization point.
+  synccpu = 0.0
+  syncsim = 0.0
+
+  # Run until asked to exit.
+  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,))
+
+    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
+        ctrlnoise = np.zeros((m.nu,))
+
+    # Sleep for 1 ms or yield, to let main thread run.
+    if simulate.run != 0 and simulate.busywait != 0:
+      time.sleep(0)
+    else:
+      time.sleep(0.001)
+
+    with simulate.lock():
+      if m is not None:
+        assert d is not None
+        if simulate.run:
+          # Record CPU time at start of iteration.
+          startcpu = glfw.get_time()
+
+          elapsedcpu = startcpu - synccpu
+          elapsedsim = d.time - syncsim
+
+          # Inject noise.
+          if simulate.ctrlnoisestd != 0.0:
+            # Convert rate and scale to discrete time (Ornstein–Uhlenbeck).
+            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]
+
+          # Requested slow-down factor.
+          slowdown = 100 / PERCENT_REALTIME[simulate.real_time_index]
+
+          # Misalignment: distance from target sim time > MAX_SYNC_MISALIGN.
+          misaligned = abs(elapsedcpu / slowdown -
+                           elapsedsim) > MAX_SYNC_MISALIGN
+
+          # Out-of-sync (for any reason): reset sync times, step.
+          if (elapsedsim < 0 or elapsedcpu < 0 or synccpu == 0 or misaligned or
+              simulate.speed_changed):
+            # Re-sync.
+            synccpu = startcpu
+            syncsim = d.time
+            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: step until ahead of cpu.
+          else:
+            measured = False
+            prevsim = d.time
+            refreshtime = SIM_REFRESH_FRACTION / simulate.refresh_rate
+            # Step while sim lags behind CPU and within refreshtime.
+            while (((d.time - syncsim) * slowdown <
+                    (glfw.get_time() - synccpu)) and
+                   ((glfw.get_time() - startcpu) < refreshtime)):
+              # Measure slowdown before first step.
+              if not measured and elapsedsim:
+                simulate.measured_slowdown = elapsedcpu / elapsedsim
+                measured = True
+
+              # Clear old perturbations, apply new.
+              d.xfrc_applied[:, :] = 0
+              simulate.applyposepertubations(0)  # Move mocap bodies only.
+              simulate.applyforceperturbations()
+
+              # Call mj_step.
+              mujoco.mj_step(m, d)
+
+              # Break if reset.
+              if d.time < prevsim:
+                break
+        else:  # simulate.run is False: GUI is paused.
+          # 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)
+
+
+def launch(model: Optional[mujoco.MjModel] = None,
+           data: Optional[mujoco.MjData] = None,
+           *,
+           run_physics_thread: bool = True) -> None:
+  """Launches the Simulate GUI."""
+  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)
+
+  # The simulate object encapsulates the UI.
+  simulate = Simulate()
+
+  # Initialize GLFW.
+  if not glfw.init():
+    raise mujoco.FatalError('could not initialize GLFW')
+
+  atexit.register(glfw.terminate)
+
+  if run_physics_thread:
+    physics_thread = threading.Thread(
+        target=_physics_loop, args=(simulate, model, data))
+    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_repl(model: mujoco.MjModel, data: mujoco.MjData) -> None:
+  """EXPERIMENTAL FEATURE: Launches the Simulate GUI in REPL mode."""
+  if typing.TYPE_CHECKING:
+    launch(model, data, run_physics_thread=False)
+    return
+
+  try:
+    import IPython  # pylint: disable=g-import-not-at-top
+    has_ipython = True
+  except ImportError:
+    has_ipython = False
+
+  def start_shell(global_variables):
+    if has_ipython and IPython.get_ipython() is not None:
+      locals().update(global_variables)
+      IPython.embed()
+    else:
+      code.InteractiveConsole(locals=global_variables).interact()
+
+  repl_thread = threading.Thread(
+      target=start_shell, args=(inspect.stack()[1][0].f_globals,))
+  repl_thread.start()
+  launch(model, data, run_physics_thread=False)
+  repl_thread.join()
+
+
+if __name__ == '__main__':
+  from absl import app  # pylint: disable=g-import-not-at-top
+  from absl import flags  # pylint: disable=g-import-not-at-top
+
+  _MJCF_PATH = flags.DEFINE_string('mjcf', None, 'Path to MJCF file.')
+
+  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)
+    else:
+      launch()
+
+  app.run(main)

python/setup.py

@@ -70,7 +70,6 @@ def get_external_lib_patterns():
   else:
     return ['libmujoco.so.*']
 
-
 def get_plugin_lib_patterns():
   if platform.system() == 'Windows':
     return ['*.dll']
@@ -322,6 +321,7 @@ setup(
         CMakeExtension('mujoco._functions'),
         CMakeExtension('mujoco._render'),
         CMakeExtension('mujoco._rollout'),
+        CMakeExtension('mujoco._simulate'),
         CMakeExtension('mujoco._structs'),
     ],
     python_requires='>=3.7',