[MISC] Avoid forcing copy in getters. (#2030)

* Remove fragile link data caching mechanism now that zero-copy is available.
* Avoid forcing copy in getters.
* Disable 'ti.u1' on Vulkan because it is broken.

Author: duburcqa

genesis/__init__.py

@@ -171,9 +171,9 @@ def init(
     tc_int = torch.int32
 
     # Bool
-    # Note that `ti.u1` is broken on Apple Metal and output garbage.
+    # Note that `ti.u1` is broken on Apple Metal and Vulkan.
     global ti_bool, np_bool, tc_bool
-    if backend == gs_backend.metal:
+    if backend in (gs_backend.metal, gs_backend.vulkan):
         ti_bool = ti.i32
         np_bool = np.int32
         tc_bool = torch.int32

genesis/engine/solvers/rigid/rigid_solver_decomp.py

@@ -14,15 +14,7 @@
 from genesis.engine.states.solvers import RigidSolverState
 from genesis.options.solvers import RigidOptions
 from genesis.utils import linalg as lu
-from genesis.utils.misc import (
-    ALLOCATE_TENSOR_WARNING,
-    DeprecationError,
-    ti_to_torch,
-    ti_to_numpy,
-    ti_to_python,
-    indices_to_mask,
-    _get_ti_metadata,
-)
+from genesis.utils.misc import ALLOCATE_TENSOR_WARNING, DeprecationError, ti_to_torch, ti_to_numpy, indices_to_mask
 from genesis.utils.sdf_decomp import SDF
 
 from ..base_solver import Solver
@@ -131,7 +123,6 @@ def __init__(self, scene: "Scene", sim: "Simulator", options: RigidOptions) -> N
         self._options = options
 
         self._cur_step = -1
-        self._links_state_cache = {}
 
         self.qpos: ti.Template | ti.types.NDArray | None = None
 
@@ -819,38 +810,10 @@ def _init_constraint_solver(self):
             else:
                 self.constraint_solver = ConstraintSolver(self)
 
-    def _get_links_data(
-        self,
-        field_name: str,
-        row_mask: slice | int | range | list | torch.Tensor | np.ndarray | None = None,
-        col_mask: slice | int | range | list | torch.Tensor | np.ndarray | None = None,
-        keepdim=True,
-        *,
-        to_torch=True,
-    ):
-        links_state_py = self._links_state_cache.setdefault((to_torch,), {})
-
-        field = getattr(self.links_state, field_name)
-        tensor = links_state_py.get(field_name)
-        if tensor is None:
-            tensor = links_state_py[field_name] = ti_to_python(field, transpose=True, to_torch=to_torch)
-
-        ti_data_meta = _get_ti_metadata(field)
-        if len(ti_data_meta.shape) < 2:
-            if row_mask is not None and col_mask is not None:
-                gs.raise_exception("Cannot specify both row and colum masks for tensor with 1D batch.")
-            mask = indices_to_mask(row_mask if col_mask is None else col_mask, keepdim=keepdim, to_torch=to_torch)
-        else:
-            mask = indices_to_mask(row_mask, col_mask, keepdim=keepdim, to_torch=to_torch)
-
-        return tensor[mask]
-
     def substep(self):
         # from genesis.utils.tools import create_timer
         from genesis.engine.couplers import SAPCoupler
 
-        self._links_state_cache.clear()
-
         kernel_step_1(
             links_state=self.links_state,
             links_info=self.links_info,
@@ -1301,7 +1264,6 @@ def set_state(self, f, state, envs_idx=None):
             self.collider.clear(envs_idx)
             if self.constraint_solver is not None:
                 self.constraint_solver.reset(envs_idx)
-            self._links_state_cache.clear()
             self._cur_step = -1
 
     def process_input(self, in_backward=False):
@@ -1541,7 +1503,6 @@ def set_base_links_pos(self, pos, links_idx=None, envs_idx=None, *, relative=Fal
             static_rigid_sim_config=self._static_rigid_sim_config,
         )
 
-        self._links_state_cache.clear()
         kernel_forward_kinematics_links_geoms(
             envs_idx,
             links_state=self.links_state,
@@ -1590,7 +1551,6 @@ def set_base_links_quat(self, quat, links_idx=None, envs_idx=None, *, relative=F
             static_rigid_sim_config=self._static_rigid_sim_config,
         )
 
-        self._links_state_cache.clear()
         kernel_forward_kinematics_links_geoms(
             envs_idx,
             links_state=self.links_state,
@@ -1668,7 +1628,6 @@ def set_qpos(self, qpos, qs_idx=None, envs_idx=None, *, skip_forward=False, unsa
                 qpos = qpos.unsqueeze(0)
             kernel_set_qpos(qpos, qs_idx, envs_idx, self._rigid_global_info, self._static_rigid_sim_config)
 
-        self._links_state_cache.clear()
         self.collider.reset(envs_idx, cache_only=True)
         if not isinstance(envs_idx, torch.Tensor):
             envs_idx = self._scene._sanitize_envs_idx(envs_idx, unsafe=unsafe)
@@ -1877,7 +1836,6 @@ def set_dofs_velocity(self, velocity, dofs_idx=None, envs_idx=None, *, skip_forw
                     velocity = velocity.unsqueeze(0)
                 kernel_set_dofs_velocity(velocity, dofs_idx, envs_idx, self.dofs_state, self._static_rigid_sim_config)
 
-        self._links_state_cache.clear()
         if not skip_forward:
             kernel_forward_velocity(
                 envs_idx,
@@ -1908,7 +1866,6 @@ def set_dofs_position(self, position, dofs_idx=None, envs_idx=None, *, unsafe=Fa
             self._static_rigid_sim_config,
         )
 
-        self._links_state_cache.clear()
         self.collider.reset(envs_idx, cache_only=True)
         self.collider.clear(envs_idx)
         if self.constraint_solver is not None:
@@ -2055,20 +2012,20 @@ def get_links_pos(
     ):
         ref = self._convert_ref_to_idx(ref)
         if ref == 0:
-            tensor = self._get_links_data("root_COM", envs_idx, links_idx, to_torch=to_torch)
+            tensor = ti_to_torch(self.links_state.root_COM, envs_idx, links_idx, transpose=True)
         elif ref == 1:
-            i_pos = self._get_links_data("i_pos", envs_idx, links_idx, to_torch=to_torch)
-            root_COM = self._get_links_data("root_COM", envs_idx, links_idx, to_torch=to_torch)
+            i_pos = ti_to_torch(self.links_state.i_pos, envs_idx, links_idx, transpose=True)
+            root_COM = ti_to_torch(self.links_state.root_COM, envs_idx, links_idx, transpose=True)
             tensor = i_pos + root_COM
         elif ref == 2:
-            tensor = self._get_links_data("pos", envs_idx, links_idx, to_torch=to_torch)
+            tensor = ti_to_torch(self.links_state.pos, envs_idx, links_idx, transpose=True)
         else:
             gs.raise_exception("'ref' must be either 'link_origin', 'link_com', or 'root_com'.")
 
         return tensor[0] if self.n_envs == 0 else tensor
 
     def get_links_quat(self, links_idx=None, envs_idx=None, *, to_torch=True, unsafe=False):
-        tensor = self._get_links_data("quat", envs_idx, links_idx, to_torch=to_torch)
+        tensor = ti_to_torch(self.links_state.quat, envs_idx, links_idx, transpose=True)
         return tensor[0] if self.n_envs == 0 else tensor
 
     def get_links_vel(
@@ -2103,7 +2060,7 @@ def get_links_vel(
         return _tensor
 
     def get_links_ang(self, links_idx=None, envs_idx=None, *, to_torch=True, unsafe=False):
-        tensor = self._get_links_data("cd_ang", envs_idx, links_idx, to_torch=to_torch)
+        tensor = ti_to_torch(self.links_state.cd_ang, envs_idx, links_idx, transpose=True)
         return tensor[0] if self.n_envs == 0 else tensor
 
     def get_links_acc(self, links_idx=None, envs_idx=None, *, unsafe=False):
@@ -2121,7 +2078,7 @@ def get_links_acc(self, links_idx=None, envs_idx=None, *, unsafe=False):
         return _tensor
 
     def get_links_acc_ang(self, links_idx=None, envs_idx=None, *, to_torch=True, unsafe=False):
-        tensor = self._get_links_data("cacc_ang", envs_idx, links_idx, to_torch=to_torch)
+        tensor = ti_to_torch(self.links_state.cacc_ang, envs_idx, links_idx, transpose=True)
         return tensor[0] if self.n_envs == 0 else tensor
 
     def get_links_root_COM(self, links_idx=None, envs_idx=None, *, to_torch=True, unsafe=False):
@@ -2131,15 +2088,15 @@ def get_links_root_COM(self, links_idx=None, envs_idx=None, *, to_torch=True, un
         This corresponds to the global COM of each entity, assuming a single-rooted structure - that is, as long as no
         two successive links are connected by a free-floating joint (ie a joint that allows all 6 degrees of freedom).
         """
-        tensor = self._get_links_data("root_COM", envs_idx, links_idx, to_torch=to_torch)
+        tensor = ti_to_torch(self.links_state.root_COM, envs_idx, links_idx, transpose=True)
         return tensor[0] if self.n_envs == 0 else tensor
 
     def get_links_mass_shift(self, links_idx=None, envs_idx=None, *, to_torch=True, unsafe=False):
-        tensor = self._get_links_data("mass_shift", envs_idx, links_idx, to_torch=to_torch)
+        tensor = ti_to_torch(self.links_state.mass_shift, envs_idx, links_idx, transpose=True)
         return tensor[0] if self.n_envs == 0 else tensor
 
     def get_links_COM_shift(self, links_idx=None, envs_idx=None, *, to_torch=True, unsafe=False):
-        tensor = self._get_links_data("i_pos_shift", envs_idx, links_idx, to_torch=to_torch)
+        tensor = ti_to_torch(self.links_state.i_pos_shift, envs_idx, links_idx, transpose=True)
         return tensor[0] if self.n_envs == 0 else tensor
 
     def get_links_inertial_mass(self, links_idx=None, envs_idx=None, *, unsafe=False):

genesis/options/renderers.py

@@ -72,7 +72,7 @@ class RayTracer(RendererOptions):
     rr_depth: int = 0
     rr_threshold: float = 0.95
 
-    # environment texure
+    # environment texture
     env_surface: Optional[Surface] = None
     env_radius: float = 1000.0
     env_pos: tuple = (0.0, 0.0, 0.0)

genesis/utils/misc.py

@@ -571,7 +571,7 @@ def _get_ti_metadata(value: ti.Field | ti.Ndarray) -> FieldMetadata:
 def ti_to_python(
     value: ti.Field | ti.Ndarray,
     transpose: bool = False,
-    copy: bool | None = True,
+    copy: bool | None = None,
     to_torch: bool = True,
 ) -> torch.Tensor | np.ndarray:
     """Converts a GsTaichi field / ndarray instance to a PyTorch tensor / Numpy array.
@@ -761,7 +761,7 @@ def ti_to_torch(
     keepdim=True,
     transpose=False,
     *,
-    copy: bool | None = True,
+    copy: bool | None = None,
 ) -> torch.Tensor:
     """Converts a GsTaichi field / ndarray instance to a PyTorch tensor.
 
@@ -800,7 +800,7 @@ def ti_to_numpy(
     keepdim=True,
     transpose=False,
     *,
-    copy: bool | None = True,
+    copy: bool | None = None,
 ) -> np.ndarray:
     """Converts a GsTaichi field / ndarray instance to a Numpy array.
 

genesis/utils/path_planning.py

@@ -83,7 +83,7 @@ def update_object(self, ee_link_idx, obj_link_idx, _pos, _quat, envs_idx):
     # ------------------------------------------------------------------------------------
 
     def _sanitize_qposs(self, qpos_goal, qpos_start, envs_idx):
-        qpos_cur = self._entity.get_qpos(envs_idx=envs_idx)
+        qpos_cur = self._entity.get_qpos(envs_idx=envs_idx).clone()
 
         qpos_goal, _, _ = self._solver._sanitize_1D_io_variables(
             qpos_goal, None, self._entity.n_qs, envs_idx, idx_name="qpos_idx", skip_allocation=True

tests/test_pbd.py

@@ -187,8 +187,8 @@ def test_cloth_attach_rigid_link(show_viewer):
     vel = np.array([[-0.0, 1.0, 0.0], [-1.0, 0.0, 0.0]], dtype=np.float32)
     box.set_dofs_velocity(vel, dofs_idx_local=[0, 1, 2])
 
-    cloth_pos0 = cloth.get_particles_pos()[:, particles_idx]
-    link_pos0 = scene.rigid_solver.links[box_link_idx].get_pos()
+    cloth_pos0 = cloth.get_particles_pos()[:, particles_idx].clone()
+    link_pos0 = scene.rigid_solver.links[box_link_idx].get_pos().clone()
 
     for _ in range(25):
         scene.step()
@@ -199,8 +199,8 @@ def test_cloth_attach_rigid_link(show_viewer):
         scene.step()
 
     # Check that the attached particles followed the link displacement per env
-    cloth_pos1 = cloth.get_particles_pos()[:, particles_idx]
-    link_pos1 = scene.rigid_solver.links[box_link_idx].get_pos()
+    cloth_pos1 = cloth.get_particles_pos()[:, particles_idx].clone()
+    link_pos1 = scene.rigid_solver.links[box_link_idx].get_pos().clone()
 
     cloth_disp = cloth_pos1 - cloth_pos0
     link_disp = link_pos1 - link_pos0
@@ -216,8 +216,8 @@ def test_cloth_attach_rigid_link(show_viewer):
         scene.step()
 
     # Make sure that the cloth is laying on the ground without moving
-    cloth_pos2 = cloth.get_particles_pos()[:, particles_idx]
-    link_pos2 = scene.rigid_solver.links[box_link_idx].get_pos()
+    cloth_pos2 = cloth.get_particles_pos()[:, particles_idx].clone()
+    link_pos2 = scene.rigid_solver.links[box_link_idx].get_pos().clone()
     cloth_disp = cloth_pos2 - cloth_pos1
     link_disp = link_pos2 - link_pos1
     link_disp = link_disp.unsqueeze(1)

tests/test_rigid_physics.py

@@ -5,6 +5,7 @@
 import tempfile
 import xml.etree.ElementTree as ET
 from contextlib import nullcontext
+from copy import deepcopy
 from typing import cast
 from pathlib import Path
 
@@ -979,17 +980,17 @@ def test_robot_scale_and_dofs_armature(xml_path, tol):
     # It is also a good opportunity to check that it updates 'invweight' and meaninertia accordingly.
     attr_orig = {}
     for scale, robot in zip(ROBOT_SCALES, scene.entities):
-        links_invweight = robot.get_links_invweight()
-        dofs_invweight = robot.get_dofs_invweight()
+        links_invweight = robot.get_links_invweight().clone()
+        dofs_invweight = robot.get_dofs_invweight().clone()
         robot.set_dofs_armature(torch.ones((robot.n_dofs,), dtype=gs.tc_float, device=gs.device))
         assert torch.all(robot.get_dofs_invweight() < 1.0)
         with pytest.raises(AssertionError):
             assert_allclose(robot.get_dofs_invweight(), dofs_invweight, tol=tol)
         with pytest.raises(AssertionError):
             assert_allclose(robot.get_links_invweight(), links_invweight, tol=tol)
         robot.set_dofs_armature(torch.zeros((robot.n_dofs,), dtype=gs.tc_float, device=gs.device))
-        links_invweight = robot.get_links_invweight()
-        dofs_invweight = robot.get_dofs_invweight()
+        links_invweight = robot.get_links_invweight().clone()
+        dofs_invweight = robot.get_dofs_invweight().clone()
         qpos = np.random.rand(robot.n_dofs)
         robot.set_dofs_position(qpos)
         robot.set_dofs_armature(torch.zeros((robot.n_dofs,), dtype=gs.tc_float, device=gs.device))
@@ -2911,7 +2912,7 @@ def must_cast(value):
 
         # Check getter and setter without row or column masking
         if getter is not None:
-            datas = getter()
+            datas = deepcopy(getter())
             is_tuple = isinstance(datas, (tuple, list))
             if arg1_max > 0:
                 assert_allclose(getter(range(arg1_max)), datas, tol=tol)
@@ -2964,7 +2965,7 @@ def must_cast(value):
                         if arg1 is None and arg2 is not None:
                             unsafe = not must_cast(arg2)
                             if getter is not None:
-                                data = getter(arg2, unsafe=unsafe)
+                                data = deepcopy(getter(arg2, unsafe=unsafe))
                             else:
                                 if is_tuple:
                                     data = [torch.ones((1, *shape)) for shape in spec]
@@ -2982,7 +2983,7 @@ def must_cast(value):
                         elif arg1 is not None and arg2 is None:
                             unsafe = not must_cast(arg1)
                             if getter is not None:
-                                data = getter(arg1, unsafe=unsafe)
+                                data = deepcopy(getter(arg1, unsafe=unsafe))
                             else:
                                 if is_tuple:
                                     data = [torch.ones((1, *shape)) for shape in spec]
@@ -3000,7 +3001,7 @@ def must_cast(value):
                         else:
                             unsafe = not any(map(must_cast, (arg1, arg2)))
                             if getter is not None:
-                                data = getter(arg1, arg2, unsafe=unsafe)
+                                data = deepcopy(getter(arg1, arg2, unsafe=unsafe))
                             else:
                                 if is_tuple:
                                     data = [torch.ones((1, 1, *shape)) for shape in spec]

tests/test_sensors.py

@@ -419,7 +419,7 @@ def test_raycaster_hits(show_viewer, n_envs):
         entity.set_pos(pos)
     if show_viewer:
         scene.visualizer.update(force=True)
-    grid_sensor_pos = grid_sensor.get_pos()
+    grid_sensor_pos = grid_sensor.get_pos().clone()
     for _ in range(100):
         scene.step()
     grid_sensor.set_pos(grid_sensor_pos)

tests/utils.py

@@ -11,6 +11,7 @@
 from functools import cache
 from itertools import chain
 from pathlib import Path
+from types import GeneratorType
 from typing import Literal, Sequence
 
 import cpuinfo
@@ -261,6 +262,8 @@ def assert_allclose(actual, desired, *, atol=None, rtol=None, tol=None, err_msg=
     # Convert input arguments as numpy arrays
     args = [actual, desired]
     for i, arg in enumerate(args):
+        if isinstance(arg, (GeneratorType, map)):
+            arg = tuple(arg)
         if isinstance(arg, (tuple, list)):
             arg = np.stack([tensor_to_array(val) for val in arg], axis=0)
         args[i] = tensor_to_array(arg)