From 1532cffdf0e04a5a734de8c0f16424acf84b2604 Mon Sep 17 00:00:00 2001
From: Google DeepMind <noreply@deepmind.com>
Date: Fri, 21 Feb 2025 10:01:04 -0800
Subject: [PATCH] #mjx Update `bind` type checks.

PiperOrigin-RevId: 729566363
Change-Id: Ifb4fd3e83988140f781ac1d5db690cf6fbe70fbf
---
 mjx/mujoco/mjx/_src/support.py | 22 +++++++++++-----------
 1 file changed, 11 insertions(+), 11 deletions(-)

diff --git a/mjx/mujoco/mjx/_src/support.py b/mjx/mujoco/mjx/_src/support.py
index 3bc2bc31..e1ab1c3d 100644
--- a/mjx/mujoco/mjx/_src/support.py
+++ b/mjx/mujoco/mjx/_src/support.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 """Engine support functions."""
-from collections.abc import Sequence
+from collections.abc import Iterable, Sequence
 from typing import Optional, Tuple, Union
 
 import jax
@@ -293,10 +293,6 @@ class BindModel(object):
   def __init__(self, model: Model, specs: Sequence[mujoco.MjStruct]):
     self.model = model
     self.prefix = ''
-    try:
-      iter(specs)
-    except TypeError:
-      specs = [specs]
     ids = []
     for spec in specs:
       if isinstance(spec, mujoco.MjsBody):
@@ -383,9 +379,13 @@ class BindModel(object):
 
 
 def _bind_model(
-    self: Model, obj: Sequence[mujoco.MjStruct]
+    self: Model, obj: mujoco.MjStruct | Iterable[mujoco.MjStruct]
 ) -> BindModel:
   """Bind a Mujoco spec to an MJX Model."""
+  if isinstance(obj, mujoco.MjStruct):
+    obj = (obj,)
+  else:
+    obj = tuple(obj)
   return BindModel(self, obj)
 
 
@@ -398,10 +398,6 @@ class BindData(object):
     self.data = data
     self.model = model
     self.prefix = ''
-    try:
-      iter(specs)
-    except TypeError:
-      specs = [specs]
     ids = []
     for spec in specs:
       if isinstance(spec, mujoco.MjsBody):
@@ -530,9 +526,13 @@ class BindData(object):
 
 
 def _bind_data(
-    self: Data, model: Model, obj: Sequence[mujoco.MjStruct]
+    self: Data, model: Model, obj: mujoco.MjStruct | Iterable[mujoco.MjStruct]
 ) -> BindData:
   """Bind a Mujoco spec to an MJX Data."""
+  if isinstance(obj, mujoco.MjStruct):
+    obj = (obj,)
+  else:
+    obj = tuple(obj)
   return BindData(self, model, obj)