Repository: Genesis-Embodied-AI/Genesis
File: genesis/engine/entities/rigid_entity/rigid_entity.py (Python)

Code before:
            Whether the position to set is absolute or relative to the initial (not current!) position. Defaults to
            False.
        zero_velocity : bool, optional
            Whether to zero the velocity of all the entity's dofs. Defaults to True. This is a safety measure after a
            sudden change in entity pose.
        envs_idx : None | array_like, optional
            The indices of the environments. If None, all environments will be considered. Defaults to None.
        """
        if not unsafe:
            _pos = torch.as_tensor(pos, dtype=gs.tc_float, device=gs.device).contiguous()
            if _pos is not pos:
                gs.logger.debug(ALLOCATE_TENSOR_WARNING)
            pos = _pos
        self._solver.set_base_links_pos(
            pos.unsqueeze(-2),
            self._base_links_idx,
            envs_idx,
            relative=relative,
            unsafe=unsafe,
            skip_forward=zero_velocity,
        )
        if zero_velocity:
            self.zero_all_dofs_velocity(envs_idx, unsafe=unsafe)

    @gs.assert_built
    def get_weld_constraints(self, as_tensor: bool = True, to_torch: bool = True):
        welds = self.scene.sim.rigid_solver.get_weld_constraints(as_tensor=as_tensor, to_torch=to_torch)
        obj_a = welds["obj_a"]
        obj_b = welds["obj_b"]
        if as_tensor:
            mask = (obj_a == self.idx) | (obj_b == self.idx)
            if self._solver.n_envs > 0:
                welds["valid_mask"] = mask
            for k in ("obj_a", "obj_b"):
                welds[k] = welds[k][mask]
            if "env" in welds:
                welds["env"] = welds["env"][mask]
        else:
            obj_a_new, obj_b_new, env_new = [], [], []
            for ea, eb, env in zip(obj_a, obj_b, welds.get("env", [0] * len(obj_a))):
                keep = (ea == self.idx) | (eb == self.idx)
                obj_a_new.append(ea[keep])
                obj_b_new.append(eb[keep])
                if "env" in welds:
                    env_new.append(env[keep])
            welds["obj_a"], welds["obj_b"] = obj_a_new, obj_b_new
            if "env" in welds:
                welds["env"] = env_new
        return welds

    @gs.assert_built

Review comment (refactor):
Remove the options `as_tensor: bool = True, to_torch: bool = True`. It should be enforced in this public method. No need for this level of flexibility here. This will simplify a lot the implementation.

Code after:
        pos : array_like
            The position to set.
        relative : bool, optional
            Whether the position to set is absolute or relative to the initial (not current!) position. Defaults to
            False.
        zero_velocity : bool, optional
            Whether to zero the velocity of all the entity's dofs. Defaults to True. This is a safety measure after a
            sudden change in entity pose.
        envs_idx : None | array_like, optional
            The indices of the environments. If None, all environments will be considered. Defaults to None.
        """
        if not unsafe:
            _pos = torch.as_tensor(pos, dtype=gs.tc_float, device=gs.device).contiguous()
            if _pos is not pos:
                gs.logger.debug(ALLOCATE_TENSOR_WARNING)
            pos = _pos
        self._solver.set_base_links_pos(
            pos.unsqueeze(-2),
            self._base_links_idx,
            envs_idx,
            relative=relative,
            unsafe=unsafe,
            skip_forward=zero_velocity,
        )
        if zero_velocity:
            self.zero_all_dofs_velocity(envs_idx, unsafe=unsafe)

    @gs.assert_built
    def get_weld_constraints(self, with_entity=None, exclude_self_contact=False):
        welds = self._solver.get_weld_constraints(as_tensor=True, to_torch=True)
        obj_a = welds["obj_a"]
        obj_b = welds["obj_b"]

        # Create mask for filtering welds involving this entity
        mask = (obj_a == self.idx) | (obj_b == self.idx)

        # Additional filtering if with_entity is specified
        if with_entity is not None:
            if self.idx == with_entity.idx:
                if exclude_self_contact:
                    gs.raise_exception("`with_entity` is self but `exclude_self_contact` is True.")
                # For self-contact, keep only self-welds
                mask = mask & ((obj_a == self.idx) & (obj_b == self.idx))
            else:
                # For cross-entity, keep welds between this entity and with_entity
                mask = mask & ((obj_a == with_entity.idx) | (obj_b == with_entity.idx))

        # Apply filtering
        for k in ("obj_a", "obj_b"):
            welds[k] = welds[k][mask]