fix site scaling to get all sites from world body

Author: Abhiram824

robosuite/models/arenas/arena.py

@@ -155,6 +155,7 @@ def set_scale(self, scale: Union[float, List[float]], obj_name: str):
         scale_mjcf_model(
             obj=obj,
             asset_root=self.asset,
+            worldbody=self.worldbody,
             scale=scale,
             get_elements_func=get_elements,
             scale_slide_joints=False,  # Arena doesn't handle slide joints

robosuite/models/objects/objects.py

@@ -97,28 +97,6 @@ def get_obj(self):
         assert self._obj is not None, "Object XML tree has not been generated yet!"
         return self._obj
 
-    def set_scale(self, scale, obj=None):
-        """
-        Scales each geom, mesh, site, body, and joint ranges (for slide joints).
-        Called during initialization but can also be used externally.
-        Args:
-            scale (float or list of floats): Scale factor (1 or 3 dims)
-            obj (ET.Element): Root object to apply scaling to. Defaults to root object of model.
-        """
-        if obj is None:
-            obj = self._obj
-
-        self._scale = scale
-
-        # Use the centralized scaling utility function
-        scale_mjcf_model(
-            obj=obj,
-            asset_root=self.asset,
-            scale=scale,
-            get_elements_func=get_elements,
-            scale_slide_joints=True,
-        )
-
     def exclude_from_prefixing(self, inp):
         """
         A function that should take in either an ET.Element or its attribute (str) and return either True or False,
@@ -506,6 +484,7 @@ def set_scale(self, scale, obj=None):
         scale_mjcf_model(
             obj=obj,
             asset_root=self.asset,
+            worldbody=self.worldbody,
             scale=scale,
             get_elements_func=get_elements,
             scale_slide_joints=False,  # MujocoXMLObject doesn't handle slide joints

robosuite/utils/mjcf_utils.py

@@ -1060,7 +1060,7 @@ def scale_site_element(element, scale_array):
         element.set("size", s_size)
 
 
-def scale_mjcf_model(obj, asset_root, scale, get_elements_func, scale_slide_joints=True):
+def scale_mjcf_model(obj, asset_root, worldbody, scale, get_elements_func, scale_slide_joints=True):
     """
     Scales all elements (geoms, meshes, bodies, joints, sites) in an MJCF model.
 
@@ -1098,7 +1098,7 @@ def scale_mjcf_model(obj, asset_root, scale, get_elements_func, scale_slide_join
         scale_joint_element(elem, scale_array, scale_slide_joints)
 
     # Scale sites
-    site_pairs = get_elements_func(obj, "site")
+    site_pairs = get_elements_func(worldbody, "site")
     for (_, elem) in site_pairs:
         scale_site_element(elem, scale_array)