[MISC] Avoid code duplication and improve zero-copy performance. (Genesis-Embodied-AI#2023)

Author: duburcqa

genesis/engine/entities/rigid_entity/rigid_link.py

@@ -10,7 +10,7 @@
 from genesis.repr_base import RBC
 from genesis.utils import geom as gu
 
-from genesis.utils.misc import DeprecationError
+from genesis.utils.misc import DeprecationError, tensor_to_array
 
 from .rigid_geom import RigidGeom, RigidVisGeom, _kernel_get_free_verts, _kernel_get_fixed_verts
 
@@ -262,7 +262,7 @@ def get_pos(self, envs_idx=None):
         envs_idx : int or array of int, optional
             The indices of the environments to get the position. If None, get the position of all environments. Default is None.
         """
-        return self._solver.get_links_pos([self._idx], envs_idx).squeeze(-2)
+        return self._solver.get_links_pos(self._idx, envs_idx).squeeze(-2)
 
     @gs.assert_built
     def get_quat(self, envs_idx=None):
@@ -274,7 +274,7 @@ def get_quat(self, envs_idx=None):
         envs_idx : int or array of int, optional
             The indices of the environments to get the quaternion. If None, get the quaternion of all environments. Default is None.
         """
-        return self._solver.get_links_quat([self._idx], envs_idx).squeeze(-2)
+        return self._solver.get_links_quat(self._idx, envs_idx).squeeze(-2)
 
     @gs.assert_built
     def get_vel(self, envs_idx=None) -> torch.Tensor:
@@ -286,7 +286,7 @@ def get_vel(self, envs_idx=None) -> torch.Tensor:
         envs_idx : int or array of int, optional
             The indices of the environments to get the linear velocity. If None, get the linear velocity of all environments. Default is None.
         """
-        return self._solver.get_links_vel([self._idx], envs_idx).squeeze(-2)
+        return self._solver.get_links_vel(self._idx, envs_idx).squeeze(-2)
 
     @gs.assert_built
     def get_ang(self, envs_idx=None) -> torch.Tensor:
@@ -298,7 +298,7 @@ def get_ang(self, envs_idx=None) -> torch.Tensor:
         envs_idx : int or array of int, optional
             The indices of the environments to get the angular velocity. If None, get the angular velocity of all environments. Default is None.
         """
-        return self._solver.get_links_ang([self._idx], envs_idx).squeeze(-2)
+        return self._solver.get_links_ang(self._idx, envs_idx).squeeze(-2)
 
     @gs.assert_built
     def get_verts(self):
@@ -545,7 +545,7 @@ def invweight(self):
         The invweight of the link.
         """
         if self._invweight is None:
-            self._invweight = self._solver.get_links_invweight([self._idx]).cpu().numpy()[..., 0, :]
+            self._invweight = tensor_to_array(self._solver.get_links_invweight(self._idx))[..., 0, :]
         return self._invweight
 
     @property

genesis/engine/sensors/raycaster.py

@@ -405,9 +405,7 @@ def build(self):
         self._shared_metadata.total_n_rays += num_rays
 
         self._shared_metadata.points_to_sensor_idx = concat_with_tensor(
-            self._shared_metadata.points_to_sensor_idx,
-            [self._idx] * num_rays,
-            flatten=True,
+            self._shared_metadata.points_to_sensor_idx, [self._idx] * num_rays, flatten=True
         )
         self._shared_metadata.return_world_frame = concat_with_tensor(
             self._shared_metadata.return_world_frame, self._options.return_world_frame

genesis/engine/solvers/base_solver.py

@@ -57,7 +57,7 @@ def set_gravity(self, gravity, envs_idx=None, *, unsafe=False):
             _kernel_set_gravity_ndarray(gravity, envs_idx, self._gravity)
 
     def get_gravity(self, envs_idx=None, *, unsafe=False):
-        tensor = ti_to_torch(self._gravity, envs_idx, transpose=True, unsafe=unsafe)
+        tensor = ti_to_torch(self._gravity, envs_idx, transpose=True)
         return tensor.squeeze(0) if self.n_envs == 0 else tensor
 
     def dump_ckpt_to_numpy(self) -> dict[str, np.ndarray]: