Enable reading textures from spec in Python

PiperOrigin-RevId: 842151742
Change-Id: I477d118cc7e2385d29b395b192fa89e9facb9c3e

doc/python.rst

@@ -654,6 +654,19 @@ The :ref:`mjsMesh` object includes convenience methods for model creation with n
    mesh = spec.add_mesh(name='prism')
    mesh.make_cone(nedge=5, radius=1)
 
+Texture editing
+^^^^^^^^^^^^^^^
+The :ref:`mjsTexture` buffer option stores the texture bytes in the ``data`` attribute. This attribute is stored as
+bytes but its elements can be read and modified as ``int``\s via index access and read as Python's built-in ``bytes`` or
+``bytearray`` objects.  See `specs_test.py
+<https://github.com/google-deepmind/mujoco/blob/main/python/mujoco/specs_test.py>`__.
+
+.. code-block:: python
+
+  texture = spec.add_texture(name='texture', height=1, width=3, nchannel=3)
+  texture.data = bytes([255, 0, 0, 0, 255, 0, 0, 0, 255])  # Assign red, green and blue pixels.
+  texture.data[1] = 255  # Change the first pixel to yellow.
+
 .. _PyMJCF:
 
 Relationship to ``PyMJCF`` and ``bind``

python/mujoco/specs.cc

@@ -134,6 +134,41 @@ void DefineArray(py::module& m, const std::string& typestr) {
       }, py::keep_alive<0, 1>(), py::return_value_policy::reference_internal);
 };
 
+// Specialization for std::byte to convert to int for Python iteration
+template <>
+void DefineArray<std::byte>(py::module& m, const std::string& typestr) {
+  using Class = MjTypeVec<std::byte>;
+  py::class_<Class>(m, typestr.c_str())
+      .def(
+          py::init([](std::byte* data, int size) { return Class(data, size); }))
+      .def("__getitem__",
+           [](Class& v, int i) -> int {
+             if (i < 0 || i >= v.size) {
+               throw py::index_error("Index out of range.");
+             }
+             return static_cast<int>(v.ptr[i]);
+           })
+      .def("__setitem__",
+           [](Class& v, int i, int c) {
+             if (i < 0 || i >= v.size) {
+               throw py::index_error("Index out of range.");
+             }
+             if (c < 0 || c > 255) {
+               throw py::value_error("Value out of range [0, 255].");
+             }
+             v.ptr[i] = static_cast<std::byte>(c);
+           })
+      .def("__len__", [](Class& v) { return v.size; })
+      .def(
+          "__iter__",
+          [](Class& v) {
+            return py::make_iterator(
+                reinterpret_cast<unsigned char*>(v.ptr),
+                reinterpret_cast<unsigned char*>(v.ptr + v.size));
+          },
+          py::keep_alive<0, 1>(), py::return_value_policy::reference_internal);
+};
+
 py::list FindAllImpl(raw::MjsBody& body, mjtObj objtype, bool recursive) {
   py::list list;
   raw::MjsElement* el = mjs_firstChild(&body, objtype, recursive);

python/mujoco/specs_test.py

@@ -1036,6 +1036,29 @@ class SpecsTest(absltest.TestCase):
     texture.data = np.zeros((2, 2, 3), dtype=np.uint8).tobytes()
     spec.compile()
 
+  def test_read_texture(self):
+    spec = mujoco.MjSpec()
+    texture = spec.add_texture(name='texture', height=1, width=2, nchannel=3)
+    texture.data = bytes([1, 2, 3, 4, 5, 6])
+    read_bytes = bytes(texture.data)
+    self.assertEqual(read_bytes, bytes([1, 2, 3, 4, 5, 6]))
+
+  def test_modify_texture(self):
+    # Assign red, green and blue pixels, then make the first pixel yellow.
+    spec = mujoco.MjSpec()
+    texture = spec.add_texture(name='texture', height=1, width=3, nchannel=3)
+    texture.data = bytes([255, 0, 0, 0, 255, 0, 0, 0, 255])
+    texture.data[1] = 255
+    self.assertEqual(
+        bytes(texture.data), bytes([255, 255, 0, 0, 255, 0, 0, 0, 255])
+    )
+
+    # Assigning values outside the range [0, 255] should raise an error.
+    with self.assertRaises(ValueError):
+      texture.data[3] = 256
+    with self.assertRaises(ValueError):
+      texture.data[3] = -1
+
   def test_find_unnamed_asset(self):
     spec = mujoco.MjSpec()
     texture_file = spec.add_texture(file='file.png')