From 9fbdcf5da4e7aec75784438d8c1a0365d923dbc3 Mon Sep 17 00:00:00 2001 From: Johnny Date: Sat, 28 Sep 2024 22:12:37 -0400 Subject: [PATCH] Add function to programmatically set object scale. Currently, only the following are scaled: geom position, mesh size, body position, site position, slide joint position and slide joint range. --- robosuite/models/objects/objects.py | 113 ++++++++++++++++++++++++++++ 1 file changed, 113 insertions(+) diff --git a/robosuite/models/objects/objects.py b/robosuite/models/objects/objects.py index 9b56dfd9cf..7f86c1bd11 100644 --- a/robosuite/models/objects/objects.py +++ b/robosuite/models/objects/objects.py @@ -452,6 +452,119 @@ def _get_geoms(self, root, _parent=None): # Return all found pairs return geom_pairs + 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 + + # Ensure scale is an array of length 3 + scale_array = np.array(self._scale).flatten() + if scale_array.size == 1: + scale_array = np.repeat(scale_array, 3) + elif scale_array.size != 3: + raise ValueError("Scale must be a scalar or a 3-element array.") + + # Scale geoms + geom_pairs = self._get_geoms(obj) + for _, (_, element) in enumerate(geom_pairs): + g_pos = element.get("pos") + g_size = element.get("size") + if g_pos is not None: + g_pos = array_to_string(string_to_array(g_pos) * scale_array) + element.set("pos", g_pos) + if g_size is not None: + g_size_np = string_to_array(g_size) + # Handle cases where size is not 3-dimensional + if len(g_size_np) == 3: + g_size_np = g_size_np * scale_array + elif len(g_size_np) == 2: + # For 2D size, assume [radius, height] for cylinders + g_size_np[0] *= np.mean(scale_array[:2]) # Average scaling in x and y + g_size_np[1] *= scale_array[2] # Scaling in z + else: + raise ValueError("Unsupported geom size dimensions.") + g_size = array_to_string(g_size_np) + element.set("size", g_size) + + # Scale meshes + meshes = self.asset.findall("mesh") + for elem in meshes: + m_scale = elem.get("scale") + if m_scale is not None: + m_scale = string_to_array(m_scale) + else: + m_scale = np.ones(3) + m_scale *= scale_array + elem.set("scale", array_to_string(m_scale)) + + # Scale bodies + body_pairs = self._get_elements(obj, "body") + for (_, elem) in body_pairs: + b_pos = elem.get("pos") + if b_pos is not None: + b_pos = string_to_array(b_pos) * scale_array + elem.set("pos", array_to_string(b_pos)) + + # Scale joints + joint_pairs = self._get_elements(obj, "joint") + for (_, elem) in joint_pairs: + j_pos = elem.get("pos") + if j_pos is not None: + j_pos = string_to_array(j_pos) * scale_array + elem.set("pos", array_to_string(j_pos)) + + # Scale joint ranges for slide joints + j_type = elem.get("type", "hinge") # Default joint type is 'hinge' if not specified + j_range = elem.get("range") + if j_range is not None and j_type == "slide": + # Get joint axis + j_axis = elem.get("axis", "1 0 0") # Default axis is [1, 0, 0] + j_axis = string_to_array(j_axis) + axis_norm = np.linalg.norm(j_axis) + if axis_norm > 0: + axis_unit = j_axis / axis_norm + else: + # Avoid division by zero + axis_unit = np.array([1.0, 0.0, 0.0]) + # Compute scaling factor along the joint axis + s = np.linalg.norm(axis_unit * scale_array) + # Scale the range + j_range_vals = string_to_array(j_range) + j_range_vals = j_range_vals * s + elem.set("range", array_to_string(j_range_vals)) + + # Scale sites + site_pairs = self._get_elements(self.worldbody, "site") + for (_, elem) in site_pairs: + s_pos = elem.get("pos") + if s_pos is not None: + s_pos = string_to_array(s_pos) * scale_array + elem.set("pos", array_to_string(s_pos)) + + s_size = elem.get("size") + if s_size is not None: + s_size_np = string_to_array(s_size) + if len(s_size_np) == 3: + s_size_np = s_size_np * scale_array + elif len(s_size_np) == 2: + s_size_np[0] *= np.mean(scale_array[:2]) # Average scaling in x and y + s_size_np[1] *= scale_array[2] # Scaling in z + elif len(s_size_np) == 1: + s_size_np *= np.mean(scale_array) + else: + raise ValueError("Unsupported site size dimensions.") + s_size = array_to_string(s_size_np) + elem.set("size", s_size) + @property def bottom_offset(self): bottom_site = self.worldbody.find("./body/site[@name='{}bottom_site']".format(self.naming_prefix))