From c8cc2d51b0b9437aaabc85d399eb3b86791756bc Mon Sep 17 00:00:00 2001
From: Google DeepMind <noreply@deepmind.com>
Date: Mon, 17 Jun 2024 06:11:17 -0700
Subject: [PATCH] Allow the viewer handle to return the mjModel/mjData objects
 from the simulation.

PiperOrigin-RevId: 643982026
Change-Id: Ida5ccb851c46f7352e724798738006be6badad0a
---
 python/mujoco/simulate.cc |  5 +++++
 python/mujoco/viewer.py   | 14 ++++++++++++++
 2 files changed, 19 insertions(+)

diff --git a/python/mujoco/simulate.cc b/python/mujoco/simulate.cc
index 35532daa..0ddd5982 100644
--- a/python/mujoco/simulate.cc
+++ b/python/mujoco/simulate.cc
@@ -123,6 +123,9 @@ class SimulateWrapper {
 
   mujoco::Simulate* simulate() { return simulate_; }
 
+  py::object GetModel() const { return m_; }
+  py::object GetData() const { return d_; }
+
  private:
   mujoco::Simulate* simulate_;
   std::atomic_int destroyed_ = 0;
@@ -222,6 +225,8 @@ PYBIND11_MODULE(_simulate, pymodule) {
       .def("lock", GetIfNotNull(&mujoco::Simulate::mtx),
            py::call_guard<py::gil_scoped_release>(),
            py::return_value_policy::reference_internal)
+      .def_property_readonly("m", &SimulateWrapper::GetModel)
+      .def_property_readonly("d", &SimulateWrapper::GetData)
       .def_property_readonly("ctrl_noise_std",
                              GetIfNotNull(&mujoco::Simulate::ctrl_noise_std),
                              py::call_guard<py::gil_scoped_release>())
diff --git a/python/mujoco/viewer.py b/python/mujoco/viewer.py
index 4bed131b..28c55e44 100644
--- a/python/mujoco/viewer.py
+++ b/python/mujoco/viewer.py
@@ -94,6 +94,20 @@ class Handle:
   def user_scn(self):
     return self._user_scn
 
+  @property
+  def m(self):
+    sim = self._sim()
+    if sim is not None:
+      return sim.m
+    return None
+
+  @property
+  def d(self):
+    sim = self._sim()
+    if sim is not None:
+      return sim.d
+    return None
+
   def close(self):
     sim = self._sim()
     if sim is not None: