From 90049343b9e39433dd6d65bf232b58e3778a5d55 Mon Sep 17 00:00:00 2001
From: Alessio Quaglino <quaglino@google.com>
Date: Thu, 9 Jan 2025 07:37:31 -0800
Subject: [PATCH] Fix texture indexing bug in bindings.

Fixes #2341.

PiperOrigin-RevId: 713667648
Change-Id: I7fa5da1506306c5374febf762ce54aff9146079b
---
 python/mujoco/bindings_test.py | 17 +++++++++++++++++
 python/mujoco/indexers.cc      |  1 +
 2 files changed, 18 insertions(+)

diff --git a/python/mujoco/bindings_test.py b/python/mujoco/bindings_test.py
index 6bfebb90..47182aa5 100644
--- a/python/mujoco/bindings_test.py
+++ b/python/mujoco/bindings_test.py
@@ -93,6 +93,19 @@ TEST_XML_PLUGIN = r"""
 </mujoco>
 """
 
+TEST_XML_TEXTURE = r"""
+<mujoco>
+  <asset>
+    <texture name="tex" type="2d" builtin="checker" rgb1=".2 .3 .4" rgb2=".1 0.15 0.2"
+      width="512" height="512" mark="cross" markrgb=".8 .8 .8"/>
+    <material name="mat" reflectance="0.3" texture="tex" texrepeat="1 1" texuniform="true"/>
+  </asset>
+  <worldbody>
+    <geom type="plane" size="1 1 1" material="mat"/>
+  </worldbody>
+</mujoco>
+"""
+
 
 @contextlib.contextmanager
 def temporary_callback(setter, callback):
@@ -1619,6 +1632,10 @@ Euler integrator, semi-implicit in velocity.
     self.assertIsNot(data1.model, data2.model)
     self.assertNotEqual(data1.model._address, data2.model._address)
 
+  def test_texture_size(self):
+    model = mujoco.MjModel.from_xml_string(TEST_XML_TEXTURE)
+    self.assertEqual(model.tex('tex').data.shape, (512, 512, 3))
+
   def _assert_attributes_equal(self, actual_obj, expected_obj, attr_to_compare):
     for name in attr_to_compare:
       actual_value = getattr(actual_obj, name)
diff --git a/python/mujoco/indexers.cc b/python/mujoco/indexers.cc
index 72311a93..158f4cae 100644
--- a/python/mujoco/indexers.cc
+++ b/python/mujoco/indexers.cc
@@ -103,6 +103,7 @@ py::array_t<T> MakeArray(T* base_ptr, int index, std::vector<int>&& shape,
     shape.insert(shape.begin(), m.hfield_nrow[index]);
   } else if (MjSize == &raw::MjModel::ntexdata) {
     offset = m.tex_adr[index];
+    shape.insert(shape.begin(), m.tex_nchannel[index]);
     shape.insert(shape.begin(), m.tex_width[index]);
     shape.insert(shape.begin(), m.tex_height[index]);
   } else if (MjSize == &raw::MjModel::nsensordata) {