From 2188cba4cd0b351d334bd084915920eea1cfc24b Mon Sep 17 00:00:00 2001
From: Nimrod Gileadi <nimrod@google.com>
Date: Thu, 20 Jun 2024 03:45:26 -0700
Subject: [PATCH] Fix a memory leak when copy.deepcopy(data) is called.

In MjDataWrapper.__deepcopy__ use the Python implementation of MjModelWrapper.__deepcopy__ to make an object that Python knows about and can release later.

Fixes google-deepmind/mujoco#1572.

PiperOrigin-RevId: 644967999
Change-Id: I95d093c672e136122e1939f69463a4d29dad82e8
---
 doc/changelog.rst                 |  4 ++
 python/mujoco/memory_leak_test.py | 69 +++++++++++++++++++++++++++++++
 python/mujoco/structs.cc          |  8 ++--
 3 files changed, 78 insertions(+), 3 deletions(-)
 create mode 100644 python/mujoco/memory_leak_test.py

diff --git a/doc/changelog.rst b/doc/changelog.rst
index cdef866c..7aa15b76 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -50,6 +50,10 @@ Simulate
 
 |br| |br| |br| |br|
 
+Python bindings
+^^^^^^^^^^^^^^^
+13. Fixed a memory leak when using ``copy.deepcopy()`` on a ``mujoco.MjData`` instance (:github:issue:`1572`).
+
 Version 3.1.6 (Jun 3, 2024)
 ---------------------------
 
diff --git a/python/mujoco/memory_leak_test.py b/python/mujoco/memory_leak_test.py
new file mode 100644
index 00000000..47fae759
--- /dev/null
+++ b/python/mujoco/memory_leak_test.py
@@ -0,0 +1,69 @@
+# Copyright 2024 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.
+# ==============================================================================
+"""Test that copying mujoco.MjData multiple times doesn't leak memory."""
+
+import copy
+import textwrap
+
+from absl.testing import absltest
+import mujoco
+
+
+class MemoryLeakTest(absltest.TestCase):
+
+  # Regression test for https://github.com/google-deepmind/mujoco/issues/1572
+  def test_deepcopy_mjdata_leak(self):
+    # MuJoCo model with textures that take up significant memory.
+    model_xml = textwrap.dedent("""
+    <mujoco>
+      <asset>
+        <material name="body"/>
+        <texture name="grid" type="2d" builtin="checker"
+            width="1024" height="1024" rgb1="0 0 0" rgb2="1 1 1"/>
+        <material name="grid" texture="grid" texrepeat="1 1"
+            texuniform="true" reflectance=".2"/>
+      </asset>
+      <worldbody>
+          <geom size="0 0 .05" type="plane" material="grid"/>
+      </worldbody>
+    </mujoco>
+    """)
+    limit = self._memory_limit(4 * 2**30)
+    try:
+      model = mujoco.MjModel.from_xml_string(model_xml)
+      data = mujoco.MjData(model)
+      total = 0
+      for _ in range(1000):
+        mujoco.mj_step(model, data)
+        last_data = copy.deepcopy(data)
+        total += last_data.time
+    finally:
+      self._memory_limit(limit)
+
+  def _memory_limit(self, limit_in_bytes: int) -> int:
+    """Limits max memory usage, and returns previous limit."""
+    soft = -1
+    try:
+      import resource  # pylint: disable=g-import-not-at-top
+      soft, hard = resource.getrlimit(resource.RLIMIT_AS)
+      resource.setrlimit(resource.RLIMIT_AS, (limit_in_bytes, hard))
+    except (ImportError, ValueError):
+      # On Windows or systems where setting resource limits fails, do nothing.
+      pass
+    return soft
+
+
+if __name__ == '__main__':
+  absltest.main()
diff --git a/python/mujoco/structs.cc b/python/mujoco/structs.cc
index a52c8335..a49d4321 100644
--- a/python/mujoco/structs.cc
+++ b/python/mujoco/structs.cc
@@ -1907,9 +1907,11 @@ This is useful for example when the MJB is not available as a file on disk.)"));
   mjData.def("__copy__", [](const MjDataWrapper& other) {
     return MjDataWrapper(other);
   });
-  mjData.def("__deepcopy__", [](const MjDataWrapper& other, py::dict) {
-    MjModelWrapper* model_copy = new MjModelWrapper(other.model());
-    return MjDataWrapper(other, model_copy);
+  mjData.def("__deepcopy__", [](const MjDataWrapper& other, py::dict memo) {
+    // Use copy.deepcopy(model) to make a model that Python is aware of.
+    py::object new_model_py =
+        py::cast(other.model()).attr("__deepcopy__")(memo);
+    return MjDataWrapper(other, new_model_py.cast<MjModelWrapper*>());
   });
   mjData.def(py::pickle(
       [](const MjDataWrapper& d) {  // __getstate__